Factor-FB-Offline — TD-MPC2 policies (Box2D)
The expert TD-MPC2 policies used to collect the Box2D split of
TommyShen/Factor-FB-Offline.
One single-task policy per Box2D task, trained from scratch on the rlmini_envs Box2D
tasks (dense potential-shaped reward; success measured from the raw sparse condition).
Metaworld / ManiSkill demos in the dataset were collected with the public multitask checkpoint
nicklashansen/newt— use that repo for those; it is not mirrored here.
Checkpoints
| task | obs_dim | act_dim | collection_success | file |
|---|---|---|---|---|
b2d-center-ctrl | 7 | 2 | 1.0 | checkpoints/box2d/b2d-center-ctrl.pt |
b2d-center-ctrl-target | 12 | 2 | 0.539 | checkpoints/box2d/b2d-center-ctrl-target.pt |
b2d-center-ctrl-target-hard | 12 | 2 | 0.49 | checkpoints/box2d/b2d-center-ctrl-target-hard.pt |
b2d-center-wall | 12 | 2 | 0.669 | checkpoints/box2d/b2d-center-wall.pt |
b2d-center-wall-navigation | 12 | 2 | 0.79 | checkpoints/box2d/b2d-center-wall-navigation.pt |
b2d-goal-ctrl | 7 | 2 | 1.0 | checkpoints/box2d/b2d-goal-ctrl.pt |
b2d-goal-ctrl-maze | 7 | 2 | 0.992 | checkpoints/box2d/b2d-goal-ctrl-maze.pt |
b2d-goal-ctrl-target | 12 | 2 | 0.691 | checkpoints/box2d/b2d-goal-ctrl-target.pt |
collection_success = mean episode success over the collection rollouts (the honest
empirical success rate; the training-time eval used only 4 episodes so was noisier).
Architecture
TD-MPC2 (model_size=B): latent_dim=512, mlp_dim=512, 5-Q ensemble, single-task
(task_dim=0). Observation/action are zero-padded to newt's multitask size
(obs→128, action→16) at inference, then stripped back to the raw dims when writing
the dataset. Files are full agent checkpoints (weights + optimizer), so they also resume
training.
Load & run (inference)
These load with newt's TD-MPC2 stack. The exact loader used to collect the data is
data_collection/collect_expert.py in
Factored-FB (branch tl-alignment):
hf download TommyShen/Factor-FB-Offline-tdmpc2 checkpoints/box2d/b2d-goal-ctrl.pt
# roll it out / collect more:
<newt>/.venv/bin/python data_collection/collect_expert.py task=b2d-goal-ctrl num_envs=50 +num_episodes=100 +success_only=true +ckpt=<downloaded>.pt +out=<dir>
Loading detail: build a model_size=B single-task WorldModel, drop the checkpoint's
_action_masks/_task_emb buffers, and load_state_dict(..., strict=False) (see
_robust_load in collect_expert.py). checkpoint_manifest.json lists per-task
obs_dim/act_dim.