chunked_policy_robomimic checkpoints

Diffusion Policy (UNet, ResNet18 + spatial softmax robomimic encoder, the paper's vision setup) trained on the robomimic ph datasets with absolute actions (*_vision_abs task configs), seed 42, batch size 64, action chunking with horizon 16, 2 observation steps and 8 executed actions.

Contents

diffusion_unet_vision/<task>_ph_seed42/
  checkpoints/epoch=XXXX-test_mean_score=Y.ckpt   # 2 best checkpoints by test success rate
  checkpoints/latest.ckpt                         # last checkpoint
  .hydra/                                         # full resolved config + command-line overrides
  logs.json.txt                                   # per-step training log

Each checkpoint (~4.6 GB) holds the raw model, the EMA model and the optimizer state, so it can be used for evaluation (use ema_model; it is the policy the rollouts evaluate) or to resume training.

Results

Test success rate over 50 rollouts (EMA policy), evaluated every 50 epochs.

task	epochs trained	gradient steps	best test success
lift	309 (stopped early)	39k	1.00
can	301	100.8k	1.00
square	601	264.4k	0.98
tool_hang	201	290.6k	0.82 (0.80 at the last epoch)

transport will be added when its training finishes.

Loading

import torch, dill, hydra

payload = torch.load("latest.ckpt", map_location="cpu", pickle_module=dill)
cfg = payload["cfg"]
workspace = hydra.utils.get_class(cfg._target_)(cfg)
workspace.load_payload(payload)
policy = workspace.ema_model   # or workspace.model
policy.eval()

This requires the chunked_policy package from the training code.