axiboai/vla-mjlab-piper-stack-pi0

VLA MJLab Piper Stack Dataset

LeRobot v3 dataset collected in MJLab for a single-arm Piper cube-stacking task:

Pick up the red cube and stack it on the blue cube.

Summary

• —Environment: vla-mjlab / MJLab bimanual Piper sim
• —Control rate / dataset FPS: 50 Hz
• —Episodes: 200
• —Successful episodes: 128
• —Policy-ready schema: 7D state / 7D action
• —Cameras: observation.images.wrist_left, observation.images.wrist_right, observation.images.scene_top
• —State: right arm q1..q6 + measured gripper position
• —Action: right arm target q1..q6 + target gripper position

The dataset is intended for ACT, Pi0/Pi05, SmolVLA, and related LeRobot policy experiments. The strongest validated baseline so far is ACT trained on the success-only subset with chunk_size=100.

Notes

This dataset is the corrected v4 collection. Compared with earlier internal collections, it fixes:

• —gripper proprioception missing from observation.state,
• —dataset FPS metadata mismatched with MJLab's 50 Hz control rate,
• —early episode cutoff before full release/retract,
• —stack success evaluation using stale/hardcoded geometry.

The blue cube is fixed and the red cube is randomized in a local workspace region. This is intended as a controlled first benchmark before broadening the task distribution.

Example ACT Training

SUCCESS_EPS=$(python - <<'EOF'
import pyarrow.parquet as pq
from pathlib import Path
root = Path.home() / "datasets/piper_stack_act_v4"
df = pq.read_table(str(root / "data/chunk-000/file-000.parquet")).to_pandas()
ok = sorted(int(e) for e, succ in df.groupby("episode_index")["metadata.success"].last().items() if succ == 1)
print(",".join(str(e) for e in ok))
EOF
)

lerobot-train \
  --policy.type=act \
  --policy.chunk_size=100 \
  --policy.n_action_steps=100 \
  --policy.n_decoder_layers=7 \
  --policy.n_encoder_layers=4 \
  --policy.dim_model=512 \
  --policy.dim_feedforward=3200 \
  --policy.kl_weight=10 \
  --dataset.repo_id=axiboai/vla-mjlab-piper-stack-act \
  --dataset.video_backend=pyav \
  --dataset.episodes="[$SUCCESS_EPS]" \
  --batch_size=8 \
  --steps=500000