Skip to content

robocasa_tabletop.md

Latest commit

 

History

History
122 lines (84 loc) · 4.09 KB

robocasa_tabletop.md

File metadata and controls

122 lines (84 loc) · 4.09 KB
Raw

RoboCasa-GR1 Tabletop: Training & Evaluation

This guide covers training and evaluation for DiT4DiT on the RoboCasa tabletop simulation benchmark.

Prepare Dataset

Download the GR00T-X simulation dataset from Hugging Face:

python examples/Robocasa_tabletop/train_files/download_gr00t_ft_data.py

This downloads 24 task folders from nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim to ./playground/Datasets/nvidia/PhysicalAI-Robotics-GR00T-X-Embodiment-Sim/.

Environment Setup

Please first follow the official RoboCasa installation guide to install the base robocasa-gr1-tabletop-tasks environment.

Then pip soceket support

pip install tyro

Configure Training

The training config is defined in DiT4DiT/config/robocasa/dit4dit_robocasa_gr1.yaml. Key parameters:

Parameter Description Default
framework.cosmos25.base_model Path to Cosmos-Predict2.5-2B -
framework.action_model.action_model_type DiT variant DiT-B
datasets.vla_data.data_root_dir Dataset root directory -
datasets.vla_data.data_mix Dataset mixture name fourier_gr1_unified_1000
datasets.vla_data.per_device_batch_size Batch size per GPU 4
num_processes Num of GPUs 16
trainer.max_train_steps Total training steps 200000
trainer.learning_rate.backbone_interface Video DiT learning rate 1e-5
trainer.learning_rate.action_model Action DiT learning rate 1e-4
trainer.freeze_modules Modules to freeze "backbone_interface.extractor.text_encoder,backbone_interface.extractor.vae"

Launch Training

Single-node:

bash examples/Robocasa_tabletop/train_files/run_robocasa.sh

Multi-node (SLURM):

sbatch examples/Robocasa_tabletop/train_files/submit_robocasa_training.sh

Note: Adjust #SBATCH -N (number of nodes) and #SBATCH --gres=gpu: (GPUs per node) in the script to control total GPU count. The total number of processes is computed automatically.

Checkpoints will be saved to {run_root_dir}/{run_id}/. Training supports:

  • DeepSpeed ZeRO Stage 2/3
  • Gradient checkpointing
  • Mixed precision (bf16)
  • Wandb logging
  • Resume from checkpoint

Inference

Download Pretrained Checkpoint

You can download our pretrained DiT4DiT-RoboCasa-GR1 checkpoint from Hugging Face to directly run evaluation:

huggingface-cli download mondo-robotics/dit4dit-model --include "dit4dit_robocasa_gr1/*" --local-dir /path/to/dit4dit-model

See the Model Zoo for all available checkpoints.

Note: After downloading, remember to update line 58 of config.yaml in the checkpoint directory to point to your local Cosmos-Predict2.5-2B path.

Option A: Single Evaluation

Step 1: Start the policy server

CUDA_VISIBLE_DEVICES=0 python deployment/model_server/server_policy.py \
  --ckpt_path /path/to/checkpoint.pt \
  --port 6398 \
  --use_bf16

Note: We use Nvidia A100 GPUs to evaluate. If you are using RTX series GPUs for evaluation, remove the --use_bf16 flag when launching server_policy.py.

Step 2: Run evaluation against the server

python examples/Robocasa_tabletop/eval_files/simulation_env.py \
  --args.env_name "gr1_unified/PnPMilkToMicrowaveClose_GR1ArmsAndWaistFourierHands_Env" \
  --args.port 6398 \
  --args.pretrained_path /path/to/checkpoint.pt \
  --args.n_episodes 50

Option B: Batch Evaluation (Multi-GPU, recommended)

Run all 24 evaluation environments across multiple GPUs:

bash examples/Robocasa_tabletop/eval_files/batch_eval_args.sh \
  /path/to/checkpoint.pt \   # Checkpoint path
  1 \                         # Number of parallel envs
  720 \                       # Max episode steps
  12 \                        # Action chunk length
  "0,1,2,3"                   # GPU IDs

This script automatically:

  1. Launches a policy server on each GPU
  2. Distributes environments across GPUs
  3. Runs 50 episodes per environment
  4. Saves videos and logs