shujialiu/robotwin3task-cosmos-p25-2b-iter36000

RoboTwin three-task Cosmos-Predict2.5-2B checkpoint (iter 36000)

This repository contains the full distributed training checkpoint for the p25_2b_224_seed0 run. It includes model, optimizer, scheduler, and trainer state and can be used to resume training rather than only for inference.

Tasks

• hanging_mug
• stack_blocks_three
• lift_pot

Checkpoint status

• Saved iteration: 36000
• Original target: 40000
• Last transient log iteration before the interrupted run: 36315
• Format: PyTorch Distributed Checkpoint (8 shards)
• Training hardware: 8 x A800 80 GB
• Micro batch per GPU: 32
• Gradient accumulation: 1
• Effective global batch: 256
• Image size: 224
• Action chunk size: 32
• Precision: bfloat16
• LoRA rank / alpha: 32 / 32

The checkpoint has not yet been validated with RoboTwin closed-loop demo_clean or demo_randomized evaluation.

Download

mkdir -p /shujia/roboooiclr
hf download shujialiu/robotwin3task-cosmos-p25-2b-iter36000 \
  --local-dir /shujia/roboooiclr/robotwin3task-cosmos-p25-2b-iter36000

The repository root is also the checkpoint path: it directly contains model/, optim/, scheduler/, and trainer/.

Resume training

Install NVIDIA Cosmos-Predict2.5 and prepare the processed RoboTwin dataset, Cosmos-Predict2.5-2B model assets, and RoboTwin simulator on the target host. Override the machine-specific paths before launching:

export ROBOTWIN3TASK_ROOT=/shujia/roboooiclr/robotwin3task-cosmos-p25-2b-iter36000
export COSMOS_ROOT=/path/to/cosmos-predict2.5
export COSMOS_MODEL_ROOT=/path/to/Cosmos-Predict2.5-2B
export ROBOTWIN_ROOT=/path/to/RoboTwin
export COSMOS_PYTHON="$COSMOS_ROOT/.venv/bin/python"
export COSMOS_TORCHRUN="$COSMOS_ROOT/.venv/bin/torchrun"

bash "$ROBOTWIN3TASK_ROOT/scripts/resume_training_highmem.sh" \
  "$ROBOTWIN3TASK_ROOT"

The resume script enforces the original effective global batch of 256. The processed demonstrations and reason1_embeddings.pkl are not included in this model repository.

Included project files

• scripts/: training, resume, export, and evaluation launchers
• configs/: evaluation protocol and seed lists
• src/: three-task dataset/config implementation
• policy/: RoboTwin policy deployment adapter
• tools/: preprocessing, audit, selection, and server utilities
• training/run/: frozen run configuration and launch metadata