CIS 6280 Stacking World Models
Pretrained world models, policies, evaluation reports, and the code snapshot for the CIS 6280 four-box stacking assignment. The matching demonstrations, cached latents, on-policy rollouts, and subgoal images are in YongYong/cis6280-stack-datasets-v2.
Contents
| Path | Description |
|---|---|
jepa/smwm-c3/model.pt |
Released world model for planning (Stage 4). Encodes each camera frame into a 192-D latent and predicts future latents from the last 3 frames and the actions in between; finetuned on 16-block on-policy rollouts of the released policy. |
jepa/smwm-c3/pretrain.pt |
Stage 2 world model before on-policy finetuning. Shares the same frozen image encoder as model.pt. |
jepa/smwm-c3/onpolicy-train-log.jsonl |
Training and validation log from Stage 4 finetuning. |
policy/df-smwm-c1-sigma04/policy.pt |
Released action-chunk policy (Stage 3). A Diffusion Forcing policy that proposes 16 action blocks (8 s) from the current frame, trained on demonstrations with action noise $\sigma \le 0.04$. |
policy/df-smwm-c1/policy.pt |
Same policy architecture trained on all demonstrations regardless of $\sigma$. |
bc/step-050000.pt |
Privileged state-based behavior-cloning teacher used to collect the 20,000-episode dataset, along with its evaluation reports. |
eval/planning-val300/ |
Per-episode benchmark reports across 300 validation scenes and 3 policy seeds (0, 2, 3), produced by assignment/training/planning/evaluate.py (commit recorded in each report's code field). |
eval/ (other folders) |
Earlier 100-scene evaluation logs using pretrain.pt. |
code/world-model-gym.tar |
Repository snapshot (including the Sharpa hand assets) used by assignment/planning/colab.ipynb; the commit is recorded in code/SNAPSHOT.txt. |
Loading Checkpoints Locally
All .pt files are torch.save dictionaries (torch.load(..., weights_only=False)). Because Stage 4 freezes the image encoder, pretrain.pt, model.pt, the policies, and the pre-encoded latents in the dataset repository all share the same 192-D latent space.
When evaluating policy.pt locally, pass --world-model path/to/jepa/smwm-c3/model.pt on the CLI (or world_model=... in Python), or update ckpt["info"]["encoder"] in the policy file as done in colab.ipynb.
Benchmark (300 Validation Scenes, 3 Policy Seeds)
Evaluated with policy/df-smwm-c1-sigma04/policy.pt, jepa/smwm-c3/model.pt, and goals/scripted-val300-every80/goals.npz over a 1795-step horizon:
| Selection | ≥1 | ≥2 | ≥3 | 4 | Grasp | Mean boxes |
|---|---|---|---|---|---|---|
one sample, no goal (LatentPolicyAgent) |
81.4% | 35.7% | 8.6% | 1.1% | 99.4% | 1.27 |
16 samples, subgoals every 80 steps, τ = 0.25 (GoalPlanningAgent) |
85.2% | 41.4% | 11.8% | 2.4% | 99.1% | 1.41 |
≥1 to 4 are the percentages of episodes whose connected stack in the goal zone reaches at least that many boxes at some point (boxes standing side by side do not count as a stack), Grasp is the percentage of episodes in which at least two hand parts press a box, and Mean boxes is the average peak number of stacked boxes. Planning stacks 0.14 more boxes per episode on average (95% confidence interval 0.06–0.22) and reaches ≥2 boxes 5.8 percentage points more often (1.3–10.1).