Skip to content

torch_world.planning

CEM

CEM(
    horizon: int,
    action_dim: int,
    num_samples: int = 300,
    num_elites: int = 30,
    iterations: int = 10,
    std: float = 1.0,
    bounds: tuple[float, float] | None = None,
    objective: str = "final",
)

Cross-entropy method over action sequences in the latent space of a WorldModel.

Each iteration samples num_samples sequences from a diagonal Gaussian, rolls them out with the model and refits the Gaussian to the num_elites cheapest under model.cost. The "final" objective scores the last imagined latent; "path" averages the cost of every imagined latent, which rewards getting close early. Under MPC, where only the first action runs, "path" keeps the planner from postponing its moves.