Factor-FB-Offline-tdmpc2

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

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.