MimicGen UR5e pi0.5 with 6-D end-effector state
This repository contains two pi0.5 checkpoints trained on the same balanced UR5e MimicGen four-task subset. The policy state contains only end-effector position and axis-angle rotation:
[ee_x, ee_y, ee_z, ee_rx, ee_ry, ee_rz]
The original gripper joint state is removed before normalization and prompt tokenization. The OSC_POSE action remains 7-D:
[delta_x, delta_y, delta_z, delta_rx, delta_ry, delta_rz, gripper_command]
Checkpoints
| Hub path | OpenPI config | Normalization |
|---|---|---|
ownnorm/14999 |
pi05_mimicgen_ur5e4_ee6_state_ownnorm_2850_15k |
UR5e-only state and action quantiles |
sharednorm/14999 |
pi05_mimicgen_ur5e4_ee6_state_sharednorm_2850_15k |
Shared Panda+UR5e state and action quantiles |
Each checkpoint includes inference parameters, normalization assets, and Orbax checkpoint metadata. Optimizer/train state is intentionally omitted.
Training setup
- Tasks: Square 950, Threading 950, Stack 475, Stack Three 475
- Total demonstrations: 2,850
- State dimensions used by the policy: 6
- Action dimensions before OpenPI padding: 7
- Action horizon: 10
- Optimizer updates: 15,000; checkpoint step: 14,999
- Micro-batch size: 8
- Gradient accumulation: 8 (effective batch size 64)
- Initialization: pi0.5 base checkpoint
OpenPI loading
Use the matching config from the DART/OpenPI codebase. For example, after
downloading ownnorm/14999:
python scripts/serve_policy.py \
--env LIBERO \
--port 8000 \
policy:checkpoint \
--policy.config pi05_mimicgen_ur5e4_ee6_state_ownnorm_2850_15k \
--policy.dir /path/to/ownnorm/14999
The environment may provide the original padded 12-D MimicGen state. The
matching policy config applies TruncateState(state_dim=6) before
normalization and pi0.5 discrete-state tokenization, so only the six
end-effector pose values affect the state tokens.