kinder-mbrl-checkpoints

kinder-mbrl — Learned World Models for KinDER

Trained world-model checkpoints for the KinDER physical-reasoning benchmark (RSS 2026).

These checkpoints are produced by the kinder-mbrl baseline repository and are intended to be used with its random-shooting MPC planner.

---

Model architecture

Each checkpoint (.pt) contains a two-head MLPDynamics network together with all normalizer statistics needed for inference.

Network topology

input: [s_t (state_dim) | a_t (action_dim)]
           │
    ┌──────▼──────┐
    │  Linear → SiLU   │  hidden_dim = 256
    │  Linear → SiLU   │  hidden_dim = 256
    └──────┬──────┘
           │ shared trunk
     ┌─────┴──────┐
     ▼            ▼
robot_head      env_head
(robot_dim)    (env_dim)
delta_robot    delta_env

Prediction rule:

delta_robot, delta_env = model(normalize(s_t), normalize(a_t))
s_{t+1} = s_t + concat(denormalize(delta_robot), denormalize(delta_env))

The robot-state and environment-state deltas are predicted by two independent linear heads sharing the same trunk. Each head has its own per-feature zero-mean unit-variance normalizer fitted on the training dataset, so the two dynamics regimes (actuator / kinematic changes vs. object / scene changes) are learned independently.

Normalizers

Four independent normalizers are stored in each checkpoint:

---

Checkpoint format

Each .pt file is a torch.save dict with the following keys:

---

Training

Models were trained from HDF5 demonstration datasets available at kinder-bench/kinder-datasets.

# Install the training package
git clone https://github.com/Princeton-Robot-Planning-and-Learning/kinder-baselines
cd kinder-baselines/kinder-mbrl
uv pip install -e ".[develop]"

# Train (1 000 epochs, batch 512, lr 1e-3 by default)
python experiments/train_world_model.py \
    --mode train \
    --hdf5_path /path/to/dataset.hdf5 \
    --output_dir output \
    --epochs 1000

# Evaluate open-loop rollout error
python experiments/train_world_model.py \
    --mode eval \
    --hdf5_path /path/to/dataset.hdf5 \
    --checkpoint output/wm.pt

Default hyperparameters:

---

Inference / MPC usage

import torch
from kinder_mbrl.planning import load_world_model, wm_get_next_state

# Load checkpoint
model, norms = load_world_model("path/to/wm.pt")

# One-step prediction
next_state = wm_get_next_state(current_state, action, model, norms)

To run the full random-shooting MPC loop:

# World-model-based planning
python experiments/run_mpc.py \
    --use_world_model \
    --checkpoint output/wm.pt \
    --num_candidates 50 \
    --horizon 5

---

Citation

If you use these datasets, please cite the paper: KinDER: A Physical Reasoning Benchmark for Robot Learning and Planning:

@inproceedings{huang2026kinder,
  title     = {KinDER: A Physical Reasoning Benchmark for Robot Learning and Planning},
  author    = {Huang, Yixuan and Li, Bowen and Saxena, Vaibhav and Liang, Yichao and Mishra, Utkarsh and Ji, Liang and Zha, Lihan and Wu, Jimmy and Kumar, Nishanth and Scherer, Sebastian and Xu, Danfei and Silver, Tom},
  booktitle = {Robotics: Science and Systems (RSS)},
  year      = {2026}
}