Name	Name	Last commit message	Last commit date
parent directory ..
launch_scripts	launch_scripts
lerobot	lerobot
olmo	olmo
scripts	scripts
tests	tests
.gitignore	.gitignore
Dockerfile	Dockerfile
LICENSE	LICENSE
README.md	README.md
pyproject.toml	pyproject.toml

MolmoAct2 Experiments

This directory contains the open-sourced MolmoAct2 training and evaluation code used to replicate our experiments and fine-tune on new LeRobot datasets.

The experiments package is organized for release use: training commands are written directly below, Beaker specs and local shell wrappers are not part of the tracked source tree, and deployment documentation lives in the vendored LeRobot docs.

Setup

Install the repository in editable mode:

git clone https://github.com/allenai/molmoact2.git
cd molmoact2/experiments
pip install -e ".[all]"
pip install -e "./lerobot[async,libero]"

Use ./lerobot[all] instead of ./lerobot[async,libero] when you need the full LeRobot hardware and simulator dependency set.

Configure local paths for your machine or scheduler:

export LEROBOT_DATA_ROOT=/path/to/lerobot/data
export LEROBOT_DEPTH_DATA_ROOT=/path/to/lerobot/depth_data
export MOLMO_DATA_DIR=/path/to/molmo/data
export SPATIAL_DATA_HOME=/path/to/spatial/embodied_training_data
export HF_HOME=/path/to/huggingface/cache
export LEROBOT_VIDEO_BACKEND=pyav
export PYTHONPATH="$PWD:$PWD/lerobot/src:${PYTHONPATH:-}"

Set credentials through environment variables or scheduler secrets:

export HF_ACCESS_TOKEN="${HF_ACCESS_TOKEN:-}"
export WANDB_API_KEY="${WANDB_API_KEY:-}"

Checkpoints

launch_scripts/train_lerobot.py accepts local checkpoint paths, URLs, and Hugging Face model IDs.

Stage	Start checkpoint
ER training	`https://storage.googleapis.com/oe-training-public/Molmo2-1225/Molmo2-4B.tar`
MolmoAct2 pretraining	`allenai/Molmo2-ER`
MolmoAct2 post-training	`allenai/MolmoAct2-Pretrain`
Standard fine-tuning	`allenai/MolmoAct2`
Depth fine-tuning	`allenai/MolmoAct2-Think`
Bimanual YAM fine-tuning	`allenai/MolmoAct2-BimanualYAM`
DROID fine-tuning	`allenai/MolmoAct2-DROID`
SO100/SO101 fine-tuning	`allenai/MolmoAct2-SO100_101`

Depth Annotations

Generate depth companion datasets before depth post-training or depth fine-tuning:

python scripts/generate_depth_annotation.py \
  "${LEROBOT_DATA_ROOT}/<repo_id>" \
  --camera-key observation.images.primary

When LEROBOT_DEPTH_DATA_ROOT is set, the generator writes to ${LEROBOT_DEPTH_DATA_ROOT}/<repo_id> by default, matching the training lookup. It writes buffer_codes with 100 entries by default, matching training's default --num_depth_tokens_per_image 100.

Fine-Tuning A New Dataset

The primary use case for this repo is adapting MolmoAct2 to a new LeRobot dataset. Add a mixture in launch_scripts/data_mixtures.py, then train with full fine-tuning, LoRA, or action-expert-only tuning.

Register A Mixture

For a single LeRobot dataset, use build_single_lerobot_mixture and register the builder in MOLMOACT2_LEROBOT_MIXTURES. libero_goal is the minimal built-in example:

def build_molmoact2_libero_goal():
    return build_single_lerobot_mixture(
        name="libero_goal",
        tag="libero",
        repo_ids=["allenai/MolmoAct2-LIBERO-Dataset"],
        action_key="action",
        state_keys=["observation.state"],
        camera_keys=[
            "observation.images.image",
            "observation.images.wrist_image",
        ],
        normalize_gripper=False,
        setup_type="single franka robotic arm in libero",
        control_mode="delta end-effector pose",
        action_horizon=10,
        n_action_steps=10,
    )


MOLMOACT2_LEROBOT_MIXTURES["libero_goal"] = build_molmoact2_libero_goal

For a new embodiment, create a new tag and describe the robot/action convention literally:

def build_molmoact2_my_robot():
    return build_single_lerobot_mixture(
        name="my_robot",
        tag="my_robot",
        repo_ids=["my-org/my-lerobot-dataset"],
        action_key="action",
        state_keys=["observation.state"],
        camera_keys=[
            "observation.images.front",
            "observation.images.wrist",
        ],
        normalize_gripper=False,
        setup_type="single robot arm in my workspace",
        control_mode="delta end-effector pose",
        action_horizon=30,
        n_action_steps=30,
    )


MOLMOACT2_LEROBOT_MIXTURES["my_robot"] = build_molmoact2_my_robot

Use the same tag only for datasets that share robot semantics, action/state normalization, camera conventions, setup_type, and control_mode.

Smoke Test

For a quick single-dataset validation, disable packing and use dynamic sequence lengths:

export EXP_NAME="molmoact2-my-robot-smoke"

HF_ACCESS_TOKEN="${HF_ACCESS_TOKEN:-}" WANDB_API_KEY="${WANDB_API_KEY:-}" torchrun \
  --standalone --nproc-per-node=1 \
  launch_scripts/train_lerobot.py \
  allenai/MolmoAct2 \
  my_robot \
  --wandb.name="${EXP_NAME}" --wandb.entity=<wandb-entity> --wandb.project=<wandb-project> \
  --max_duration=20 \
  --device_batch_size=1 \
  --global_batch_size=1 \
  --num_workers=0 --pin_memory=false \
  --save_folder="checkpoints/smoke/${EXP_NAME}" \
  --packing=false \
  --dynamic_seq_len=true \
  --ft_vlm=false \
  --ft_action_expert=true \
  --ft_embedding=none

Full Fine-Tuning

Full fine-tuning updates the VLM, vision tower, connector, LM head, and action expert. Use it for larger datasets or substantial embodiment changes.

export EXP_NAME="molmoact2-my-robot-fft"

HF_ACCESS_TOKEN="${HF_ACCESS_TOKEN:-}" WANDB_API_KEY="${WANDB_API_KEY:-}" torchrun \
  --nnodes="${NNODES:-1}" --nproc-per-node=8 \
  --node_rank="${RANK:-0}" --master_addr="${ADDR:-127.0.0.1}" --master_port="${PORT:-29415}" \
  launch_scripts/train_lerobot.py \
  allenai/MolmoAct2 \
  my_robot \
  --wandb.name="${EXP_NAME}" --wandb.entity=<wandb-entity> --wandb.project=<wandb-project> \
  --max_duration=50000 \
  --device_batch_size=2 \
  --global_batch_size=64 \
  --num_workers=4 --pin_memory=true \
  --data.timeout=900 \
  --save_interval=10000 \
  --save_num_checkpoints_to_keep=20 \
  --save_folder="checkpoints/finetune/${EXP_NAME}" \
  --packing=false \
  --dynamic_seq_len=true \
  --ft_vlm=true \
  --ft_action_expert=true \
  --ft_embedding=lm_head \
  --lora_enable=false \
  --llm_learning_rate=1e-5 \
  --vit_learning_rate=5e-6 \
  --connector_learning_rate=5e-6 \
  --action_expert_learning_rate=5e-5

LoRA Fine-Tuning

LoRA fine-tuning updates LoRA adapters on the VLM path and fully trains the action expert. Use it for smaller datasets or similar embodiments.

export EXP_NAME="molmoact2-my-robot-lora"

HF_ACCESS_TOKEN="${HF_ACCESS_TOKEN:-}" WANDB_API_KEY="${WANDB_API_KEY:-}" torchrun \
  --nnodes="${NNODES:-1}" --nproc-per-node=8 \
  --node_rank="${RANK:-0}" --master_addr="${ADDR:-127.0.0.1}" --master_port="${PORT:-29415}" \
  launch_scripts/train_lerobot.py \
  allenai/MolmoAct2 \
  my_robot \
  --wandb.name="${EXP_NAME}" --wandb.entity=<wandb-entity> --wandb.project=<wandb-project> \
  --max_duration=50000 \
  --device_batch_size=2 \
  --global_batch_size=64 \
  --num_workers=4 --pin_memory=true \
  --data.timeout=900 \
  --save_interval=10000 \
  --save_num_checkpoints_to_keep=20 \
  --save_folder="checkpoints/finetune/${EXP_NAME}" \
  --packing=false \
  --dynamic_seq_len=true \
  --ft_vlm=true \
  --ft_action_expert=true \
  --ft_embedding=lm_head \
  --lora_enable=true \
  --lora_rank=64 \
  --llm_learning_rate=5e-5 \
  --vit_learning_rate=5e-5 \
  --connector_learning_rate=5e-5 \
  --action_expert_learning_rate=5e-5

Action-Expert-Only Fine-Tuning

Action-expert-only fine-tuning freezes the VLM, vision tower, connector, embeddings, and LM head. Use it when the vision-language behavior should remain fixed and the new dataset mainly changes the continuous control head.

export EXP_NAME="molmoact2-my-robot-ae-only"

HF_ACCESS_TOKEN="${HF_ACCESS_TOKEN:-}" WANDB_API_KEY="${WANDB_API_KEY:-}" torchrun \
  --nnodes="${NNODES:-1}" --nproc-per-node=8 \
  --node_rank="${RANK:-0}" --master_addr="${ADDR:-127.0.0.1}" --master_port="${PORT:-29415}" \
  launch_scripts/train_lerobot.py \
  allenai/MolmoAct2 \
  my_robot \
  --wandb.name="${EXP_NAME}" --wandb.entity=<wandb-entity> --wandb.project=<wandb-project> \
  --max_duration=50000 \
  --device_batch_size=2 \
  --global_batch_size=64 \
  --num_workers=4 --pin_memory=true \
  --data.timeout=900 \
  --save_interval=10000 \
  --save_num_checkpoints_to_keep=20 \
  --save_folder="checkpoints/finetune/${EXP_NAME}" \
  --packing=false \
  --dynamic_seq_len=true \
  --ft_vlm=false \
  --ft_action_expert=true \
  --ft_embedding=none \
  --lora_enable=false \
  --action_expert_learning_rate=5e-5

Depth Fine-Tuning

For depth-reasoning fine-tuning, start from allenai/MolmoAct2-Think and add:

--enable_depth_reasoning=true \
--num_depth_tokens=128 \
--num_depth_tokens_per_image=100 \
--depth_code_input_noise_rate=0.1 \
--style_robot_action=1.0 \
--style_robot_depth=0.0 \
--style_robot_depth_action=1.0

Application

The vendored lerobot package is the supported inference and deployment path for MolmoAct2. It covers simulator rollout, real-robot rollout/recording, async inference, CUDA graph inference, and the MolmoAct2 policy configuration surface.

See lerobot/docs/source/molmoact2.mdx for deployment commands and LeRobot-specific details.

Reproducing Released Training Stages

The commands below reproduce the major MolmoAct2 training stages. They are references for release replication; for new projects, start with the new-dataset fine-tuning section above.

Molmo2-ER Training

Molmo2-ER starts from the public Molmo2-4B checkpoint archive, not from a Hugging Face model. The launcher downloads and extracts this archive into ${MOLMOACT2_CHECKPOINT_CACHE:-${HF_HOME}/molmoact2/checkpoints} before loading it.

export EXP_NAME="molmo2-er"
export SPATIAL_DATA_HOME="${SPATIAL_DATA_HOME:-/path/to/spatial/embodied_training_data}"

WANDB_API_KEY="${WANDB_API_KEY:-}" torchrun \
  --nnodes=2 --nproc-per-node=8 \
  --node_rank="${RANK}" --master_addr="${ADDR}" --master_port="${PORT}" \
  launch_scripts/sft.py \
  https://storage.googleapis.com/oe-training-public/Molmo2-1225/Molmo2-4B.tar \
  molmo2_embodied_spatial_mix_50 \
  --seq_len=16384 \
  --device_batch_size=1 \
  --max_duration=20000 \
  --save_interval=1000 \
  --save_num_checkpoints_to_keep=20 \
  --num_workers=2 \
  --data.persistent_workers=true \
  --save_folder="checkpoints/er/${EXP_NAME}" \
  --wandb.name="${EXP_NAME}" \
  --wandb.entity=<wandb-entity> \
  --wandb.project=<wandb-project> \
  --model.vision_backbone.compile_connector=null \
  --model.mm_preprocessor.max_frames=128 \
  --model.mm_preprocessor.use_frame_special_tokens=true \
  --model.mm_preprocessor.max_subtitle_tokens=null

Use spatial-all-v4 for the spatial-only stage, or molmo2_embodied_spatial_mix_30, molmo2_embodied_spatial_mix_50, molmo2_embodied_spatial_mix_70, and molmo2_embodied_spatial_mix_90 for general/spatial recovery ablations.

MolmoAct2 Pretraining

export EXP_NAME="molmoact2-pretrain"

HF_ACCESS_TOKEN="${HF_ACCESS_TOKEN:-}" WANDB_API_KEY="${WANDB_API_KEY:-}" torchrun \
  --nnodes=8 --nproc-per-node=8 \
  --node_rank="${RANK}" --master_addr="${ADDR}" --master_port="${PORT}" \
  launch_scripts/train_lerobot.py \
  allenai/Molmo2-ER \
  pre_post_train \
  --wandb.name="${EXP_NAME}" --wandb.entity=<wandb-entity> --wandb.project=<wandb-project> \
  --seq_len=4200 \
  --max_duration=200000 \
  --device_batch_size=2 \
  --global_batch_size=128 \
  --log_interval=20 \
  --num_workers=4 --pin_memory=true \
  --data.timeout=900 \
  --save_interval=10000 \
  --save_num_checkpoints_to_keep=20 \
  --save_folder="checkpoints/pretrain/${EXP_NAME}" \
  --packing=true \
  --crop_mode=resize \
  --ft_embedding=added_tokens \
  --add_action_expert=false \
  --action_format=discrete \
  --ft_vlm=true \
  --connector_learning_rate=5e-6 \
  --vit_learning_rate=5e-6 \
  --llm_learning_rate=1e-5 \
  --random_camera_order=episode \
  --frame_loading_backend=torchcodec_exact \
  --use_annotated_task=true \
  --sample_annotated_task=false

Post-Training

export EXP_NAME="molmoact2-posttrain"

HF_ACCESS_TOKEN="${HF_ACCESS_TOKEN:-}" WANDB_API_KEY="${WANDB_API_KEY:-}" torchrun \
  --nnodes=8 --nproc-per-node=8 \
  --node_rank="${RANK}" --master_addr="${ADDR}" --master_port="${PORT}" \
  launch_scripts/train_lerobot.py \
  allenai/MolmoAct2-Pretrain \
  pre_post_train \
  --wandb.name="${EXP_NAME}" --wandb.entity=<wandb-entity> --wandb.project=<wandb-project> \
  --seq_len=2100 --vlm_seq_len=4200 \
  --max_duration=100000 \
  --device_batch_size=2 \
  --global_batch_size=128 \
  --num_workers=4 --pin_memory=true \
  --data.timeout=900 \
  --save_interval=10000 \
  --save_num_checkpoints_to_keep=20 \
  --save_folder="checkpoints/posttrain/${EXP_NAME}" \
  --separate_vlm_dataloader=true \
  --packing=true \
  --pad_packed_action_chunks=true \
  --packed_action_chunk_cap=5 \
  --crop_mode=resize \
  --ft_embedding=added_tokens \
  --ft_vlm=true \
  --ft_action_expert=true \
  --action_expert_learning_rate=5e-5 \
  --num_flow_timesteps=4 \
  --mask_action_dim_padding=true \
  --random_camera_order=episode \
  --frame_loading_backend=torchcodec_exact \
  --use_annotated_task=true \
  --sample_annotated_task=false

Depth post-training:

export EXP_NAME="molmoact2-posttrain-depth"

HF_ACCESS_TOKEN="${HF_ACCESS_TOKEN:-}" WANDB_API_KEY="${WANDB_API_KEY:-}" torchrun \
  --nnodes=8 --nproc-per-node=8 \
  --node_rank="${RANK}" --master_addr="${ADDR}" --master_port="${PORT}" \
  launch_scripts/train_lerobot.py \
  allenai/MolmoAct2-Pretrain \
  pre_post_train \
  --wandb.name="${EXP_NAME}" --wandb.entity=<wandb-entity> --wandb.project=<wandb-project> \
  --seq_len=2100 --vlm_seq_len=4200 \
  --max_duration=100000 \
  --device_batch_size=2 \
  --global_batch_size=128 \
  --num_workers=4 --pin_memory=true \
  --data.timeout=900 \
  --save_interval=10000 \
  --save_num_checkpoints_to_keep=20 \
  --save_folder="checkpoints/posttrain/${EXP_NAME}" \
  --separate_vlm_dataloader=true \
  --skip_missing_vlm_examples=true \
  --packing=true \
  --pad_packed_action_chunks=true \
  --packed_action_chunk_cap=5 \
  --crop_mode=resize \
  --ft_embedding=added_tokens \
  --ft_vlm=true \
  --ft_action_expert=true \
  --action_expert_learning_rate=5e-5 \
  --num_flow_timesteps=4 \
  --mask_action_dim_padding=true \
  --random_camera_order=episode \
  --frame_loading_backend=torchcodec_exact \
  --use_annotated_task=true \
  --sample_annotated_task=false \
  --num_depth_tokens=128 \
  --enable_depth_reasoning=true \
  --num_depth_tokens_per_image=100 \
  --style_robot_action=1.0 \
  --style_robot_depth=1.0 \
  --style_robot_depth_action=1.0

Fine-Tuning Existing Mixtures

Use the closest released checkpoint for the target embodiment:

General LIBERO or Franka tabletop tasks: allenai/MolmoAct2
Depth reasoning: allenai/MolmoAct2-Think
DROID: allenai/MolmoAct2-DROID
Bimanual YAM: allenai/MolmoAct2-BimanualYAM
SO100/SO101: allenai/MolmoAct2-SO100_101

DROID full fine-tuning:

export EXP_NAME="molmoact2-droid"

HF_ACCESS_TOKEN="${HF_ACCESS_TOKEN:-}" WANDB_API_KEY="${WANDB_API_KEY:-}" torchrun \
  --nnodes=4 --nproc-per-node=8 \
  --node_rank="${RANK}" --master_addr="${ADDR}" --master_port="${PORT}" \
  launch_scripts/train_lerobot.py \
  allenai/MolmoAct2-DROID \
  droid \
  --wandb.name="${EXP_NAME}" --wandb.entity=<wandb-entity> --wandb.project=<wandb-project> \
  --seq_len=2100 \
  --max_duration=50000 \
  --device_batch_size=2 \
  --global_batch_size=64 \
  --num_workers=4 --pin_memory=true \
  --data.timeout=900 \
  --save_interval=10000 \
  --save_num_checkpoints_to_keep=20 \
  --save_folder="checkpoints/finetune/${EXP_NAME}" \
  --packing=true \
  --crop_mode=resize \
  --ft_vlm=true \
  --ft_action_expert=true \
  --action_expert_learning_rate=5e-5 \
  --num_flow_timesteps=8 \
  --mask_action_dim_padding=true \
  --random_camera_order=none \
  --frame_loading_backend=torchcodec_exact \
  --use_annotated_task=false \
  --sample_annotated_task=false

Depth fine-tuning uses the same mixture names and --enable_depth_reasoning=true:

export EXP_NAME="molmoact2-droid-depth"

HF_ACCESS_TOKEN="${HF_ACCESS_TOKEN:-}" WANDB_API_KEY="${WANDB_API_KEY:-}" torchrun \
  --nnodes=4 --nproc-per-node=8 \
  --node_rank="${RANK}" --master_addr="${ADDR}" --master_port="${PORT}" \
  launch_scripts/train_lerobot.py \
  allenai/MolmoAct2-Think \
  droid \
  --wandb.name="${EXP_NAME}" --wandb.entity=<wandb-entity> --wandb.project=<wandb-project> \
  --seq_len=2100 \
  --max_duration=50000 \
  --device_batch_size=2 \
  --global_batch_size=64 \
  --num_workers=4 --pin_memory=true \
  --data.timeout=900 \
  --save_interval=10000 \
  --save_num_checkpoints_to_keep=20 \
  --save_folder="checkpoints/finetune/${EXP_NAME}" \
  --packing=true \
  --crop_mode=resize \
  --ft_vlm=true \
  --ft_action_expert=true \
  --action_expert_learning_rate=5e-5 \
  --num_flow_timesteps=8 \
  --mask_action_dim_padding=true \
  --random_camera_order=none \
  --frame_loading_backend=torchcodec_exact \
  --use_annotated_task=false \
  --sample_annotated_task=false \
  --num_depth_tokens=128 \
  --enable_depth_reasoning=true \
  --num_depth_tokens_per_image=100 \
  --depth_code_input_noise_rate=0.1 \
  --style_robot_action=1.0 \
  --style_robot_depth=0.0 \
  --style_robot_depth_action=1.0

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

README.md

MolmoAct2 Experiments

Setup

Checkpoints

Depth Annotations

Fine-Tuning A New Dataset

Register A Mixture

Smoke Test

Full Fine-Tuning

LoRA Fine-Tuning

Action-Expert-Only Fine-Tuning

Depth Fine-Tuning

Application

Reproducing Released Training Stages

Molmo2-ER Training

MolmoAct2 Pretraining

Post-Training

Fine-Tuning Existing Mixtures

Uh oh!

FilesExpand file tree

experiments

Directory actions

More options

Directory actions

More options

Latest commit

History

experiments

Folders and files

parent directory

README.md

MolmoAct2 Experiments

Setup

Checkpoints

Depth Annotations

Fine-Tuning A New Dataset

Register A Mixture

Smoke Test

Full Fine-Tuning

LoRA Fine-Tuning

Action-Expert-Only Fine-Tuning

Depth Fine-Tuning

Application

Reproducing Released Training Stages

Molmo2-ER Training

MolmoAct2 Pretraining

Post-Training

Fine-Tuning Existing Mixtures