Hierarchical Planning with Latent World Models

The macro-action space is trained jointly with the high-level world model through a learned action encoder that compresses variable-length sequences of primitive, low-level actions into compact latent representations.

Training process:

1. Waypoint Selection: During training, the system extracts sequences from offline trajectory data and selects specific "waypoint" states separated by variable time steps.
2. Action Encoding: The raw sequence of primitive actions executed between two consecutive waypoints is passed into the action encoder, which summarizes them into a single latent macro-action.
3. Joint Optimization via Next-Latent Prediction: These encoded macro-actions are fed into the high-level world model alongside the latent features of the starting waypoint. The high-level model's objective is to predict the latent representation of the next waypoint.

Both the action encoder and the high-level world model are trained simultaneously using a "teacher-forcing" loss function, which minimizes the error (L1 distance) between the predicted next waypoint and the actual next waypoint from the training data.

Encoder architecture:

Varies based on the environment:

• Franka arm and Push-T: For visual robotic manipulation tasks transformer-based encoder where a CLS¹ token is passed through an MLP head to output the latent action.
• For 2D maze navigation: a simple MLP architecture.

The world models in this framework are trained entirely using an offline learning paradigm on large, pre-collected datasets of task-agnostic, reward-free offline trajectories, meaning the model does not actually do any active exploration itself.

Data source:

1. Real-World Robotic Tasks (Franka Arm): The models are trained on approximately 130 hours of pre-existing, unlabeled video clips of real robots performing various manipulation tasks, taken from the massive open-source DROID and RoboSet datasets.
2. Push-T Simulation: The models use a pre-existing offline dataset consisting of 18,500 trajectories of a T-shaped block being pushed.
3. Diverse Maze Navigation: The training data consists of 5 million transitions collected beforehand by placing an agent in 25 different maze maps and simply having it execute randomly sampled actions.

Training the Low-Level World Model:

The low-level model learns the fine-grained physics of the environment by predicting the immediate next step based on a primitive action. It is trained using a combination of two loss functions:

1. Teacher-Forcing Loss: The model predicts the very next latent state, and the loss is the L1 distance (error) between its prediction and the actual next state from the offline data.
2. Multi-Step Autoregressive Rollout Loss: To prevent errors from compounding quickly, the model is also forced to predict multiple steps ahead continuously, minimizing the error of the final predicted state across a longer sequence.

Training the High-Level World Model:

The high-level model learns long-horizon dynamics by skipping timesteps and predicting the latent representation of distant "waypoints".

1. Take initial waypoint and a sequence of primitive actions
2. Compress the action sequence to a macro action
3. Predict the future waypoint
4. Use teacher-forcing loss function: minimizing the L1 distance between the predicted and actual next waypoint.