chunked_policy_robomimic checkpoints

Diffusion Policy (UNet) 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, with two vision encoders:

• diffusion_unet_vision/: ResNet18 + spatial softmax robomimic encoder (the paper's vision setup).
• diffusion_unet_vision_vit/: MinVit encoder from IBRL with IBRL's robomimic settings (embed_style: embed2, depth: 1, embed_dim: 128, num_head: 4), one ViT per camera, patch tokens flattened and projected to 64 features. Everything else is identical.

Contents

diffusion_unet_vision[_vit]/<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 with ResNet, ~4.1 GB with ViT) 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.

ResNet encoder (diffusion_unet_vision/)

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)

ViT encoder (diffusion_unet_vision_vit/)

task	epochs trained	gradient steps	best test success
lift	351	44.6k	1.00 (1.00 at the last epoch)
can	351	117.6k	0.98 (0.94 at the last epoch)
square	601	264.4k	0.86 (0.82 at the last epoch)
tool_hang	151	218.3k	0.68 (0.50 at the last epoch)
transport	151	213.7k	0.94 (0.94 at the last epoch)

ResNet 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.