Encode a 6D observation [x, y, vx, vy, goal_x, goal_y] into a latent vector.

Returns the latent embedding produced by the trained JEPA encoder. The latent is the space the world model predicts in; downstream tools operate on these latents rather than on raw observations.

Roll the learned latent dynamics forward from a latent over an action sequence.

Given a starting latent and a list of 2D acceleration actions, applies the predictor g step by step and returns the latent trajectory. This is the model imagining a future without touching the real environment.

Plan the next action that drives the agent toward a goal, via latent MPC.

Encodes the current and goal observations, searches action sequences by rolling the world model forward in latent space (cross-entropy method), and returns the first action of the best plan. Call repeatedly (replanning each step) to follow a receding-horizon trajectory to the goal.

Reset the point-mass environment to a deterministic start for the given seed.

Returns the initial observation, the agent state, and the goal position.

Apply a 2D acceleration action to the environment and advance one timestep.

Returns the new observation, reward (negative distance to goal), a done flag (agent reached the goal), and the distance to the goal.