diff --git a/docs/source/_toctree.yml b/docs/source/_toctree.yml index 13aae377b..a1e019f6f 100644 --- a/docs/source/_toctree.yml +++ b/docs/source/_toctree.yml @@ -61,6 +61,8 @@ title: π₀-FAST (Pi0Fast) - local: pi05 title: π₀.₅ (Pi05) + - local: molmoact2 + title: MolmoAct2 - local: eo1 title: EO-1 - local: groot @@ -75,6 +77,10 @@ - sections: - local: sarm title: SARM + - local: robometer + title: ROBOMETER + - local: topreward + title: TOPReward title: "Reward Models" - sections: - local: inference diff --git a/docs/source/molmoact2.mdx b/docs/source/molmoact2.mdx new file mode 100644 index 000000000..ddd178acd --- /dev/null +++ b/docs/source/molmoact2.mdx @@ -0,0 +1,433 @@ +# MolmoAct2 Policy + +MolmoAct2 is the LeRobot policy implementation of +[MolmoAct2](https://allenai.org/blog/molmoact2), ported into the LeRobot +training, evaluation, checkpointing, and dataset interfaces for easier use with +LeRobot datasets. + +This implementation currently supports training and evaluation for the regular +MolmoAct2 model. MolmoAct2-Think, which supports adaptive depth reasoning, is +not included in this LeRobot policy yet and is coming soon. + +For the original MolmoAct2 training code used for the experiments reported in +the paper, see [allenai/molmoact2](https://github.com/allenai/molmoact2). + +## Installation Requirements + +Install LeRobot with the MolmoAct2 optional dependencies: + +```bash +pip install -e ".[molmoact2]" +``` + +To run the models in this repository, you need an NVIDIA GPU. The measurements +below were taken on a single NVIDIA H100 80GB with bf16 model loading, LIBERO with two RGB cameras. MolmoAct2 rows use `chunk_size=10`, action dim 7 +padded to `expected_max_action_dim=32`, and `num_flow_timesteps=8`. Training measurements use +`gradient_checkpointing=true` and include the forward pass, backward pass, +gradient clipping, optimizer step, and optimizer state allocation. Values are +peak GPU memory sampled with `nvidia-smi`. Leave a few GiB of headroom for +dataloader workers, CUDA context, and fragmentation. + +Multi-GPU training through `accelerate` increases throughput and global batch +size, but this LeRobot port does not currently expose the original MolmoAct2 +`fsdp_devices` model-parallel training path. The current training script has +not been tested for multi-node training. + +| Mode | Peak Memory, bs=8 | Peak Memory, bs=16 | Peak Memory, bs=32 | +| ------------------------------------------------ | ----------------: | -----------------: | -----------------: | +| Inference, continuous, CUDA graph enabled (bs=1) | 12.1 GiB | - | - | +| Fine-tuning, action expert only, continuous | 16.5 GiB | 18.3 GiB | 21.4 GiB | +| Fine-tuning, LoRA VLM, both action modes | 20.2 GiB | 26.8 GiB | 41.3 GiB | +| Fine-tuning, full model, both action modes | 48.3 GiB | 49.8 GiB | 60.1 GiB | + +The repo has been tested with Ubuntu 22.04. + +## Usage + +To use MolmoAct2 in a LeRobot training config, set: + +```python +policy.type=molmoact2 +``` + +## Training + +MolmoAct2 can be fine-tuned from either the released MolmoAct2 Hugging Face +checkpoint format or from a checkpoint already saved by LeRobot. Both routes use +the same LeRobot training loop, dataset transforms, checkpoint saving, and +logging. The difference is only how the initial policy weights and processor +state are loaded. + +### Training With Original MolmoAct2 Weight + +Use `policy.checkpoint_path` when starting from a released MolmoAct2 checkpoint, +for example `allenai/MolmoAct2` or `allenai/MolmoAct2-LIBERO`. LeRobot will load +the original HF model files, then build its own policy processor from the +dataset metadata and the policy options below. + +The command below shows full fine-tuning on the merged LIBERO dataset. It uses +bf16 model loading, 8 flow timesteps, LeRobot dataset statistics, image +augmentation, and LeRobot's checkpointing/logging path. + +```bash +accelerate launch \ + --num_processes=8 \ + --mixed_precision=bf16 \ + -m lerobot.scripts.lerobot_train \ + --dataset.repo_id=allenai/MolmoAct2-LIBERO-Dataset \ + --dataset.root=/path/to/lerobot/data/allenai/MolmoAct2-LIBERO-Dataset \ + --dataset.video_backend=pyav \ + --dataset.image_transforms.enable=true \ + --policy.type=molmoact2 \ + --policy.checkpoint_path=allenai/MolmoAct2-LIBERO \ + --policy.device=cuda \ + --policy.action_mode=both \ + --policy.chunk_size=10 \ + --policy.n_action_steps=10 \ + --policy.setup_type="single franka robotic arm in libero" \ + --policy.control_mode="delta end-effector pose" \ + --policy.image_keys='["observation.images.image","observation.images.wrist_image"]' \ + --policy.model_dtype=bfloat16 \ + --policy.num_flow_timesteps=8 \ + --policy.gradient_checkpointing=true \ + --policy.freeze_embedding=true \ + --policy.normalize_gripper=false \ + --policy.enable_knowledge_insulation=false \ + --policy.push_to_hub=false \ + --wandb.enable=true \ + --wandb.entity= \ + --wandb.project= \ + --job_name= \ + --output_dir=outputs/ \ + --steps=10000 \ + --batch_size=32 \ + --num_workers=4 \ + --log_freq=20 \ + --eval_freq=-1 \ + --save_checkpoint=true \ + --save_freq=2000 +``` + +### Training With LeRobot MolmoAct2 Weight + +Use `policy.path` when starting from a MolmoAct2 checkpoint that was saved by +LeRobot, either from a local `pretrained_model` directory or from the Hub. This +restores the saved LeRobot policy config, model weights, processor, and +normalization statistics. You can still override training-time options such as +`batch_size`, `steps`, LoRA flags, or `policy.action_mode`. + +```bash +accelerate launch \ + --num_processes=8 \ + --mixed_precision=bf16 \ + -m lerobot.scripts.lerobot_train \ + --dataset.repo_id=allenai/MolmoAct2-LIBERO-Dataset \ + --dataset.root=/path/to/lerobot/data/allenai/MolmoAct2-LIBERO-Dataset \ + --dataset.video_backend=pyav \ + --dataset.image_transforms.enable=true \ + --policy.path=/path/to/pretrained_model \ + --policy.device=cuda \ + --policy.action_mode=both \ + --policy.chunk_size=10 \ + --policy.n_action_steps=10 \ + --policy.model_dtype=bfloat16 \ + --policy.num_flow_timesteps=8 \ + --policy.gradient_checkpointing=true \ + --wandb.enable=true \ + --wandb.entity= \ + --wandb.project= \ + --job_name= \ + --output_dir=outputs/ \ + --steps=10000 \ + --batch_size=32 \ + --num_workers=4 \ + --log_freq=20 \ + --eval_freq=-1 \ + --save_checkpoint=true \ + --save_freq=2000 +``` + +### Common Practices + +For fine-tuning on a comparatively small dataset, such as a single LIBERO suite +or a real-world dataset with less than 200 demonstrations, a global batch size of +16 to 32 is a good starting point. In these settings, `policy.enable_lora_vlm=true` or `policy.train_action_expert_only=true` is also a practical choice. In both +cases, we intentionally keep the action expert fully trainable, which we found +to be crucial for model performance. For larger fine-tuning datasets, larger +global batch sizes and full fine-tuning are usually preferred. + +### Common Policy Options + +- `policy.checkpoint_path`: original MolmoAct2 HF checkpoint to initialize from. + Use this for released MolmoAct2 weights. +- `policy.path`: LeRobot checkpoint to initialize from. Use this for checkpoints + created by LeRobot training. +- `policy.action_mode`: training target, one of `continuous`, `discrete`, or + `both`. `both` trains the flow-matching action expert and the discrete + action-token loss. +- `policy.train_action_expert_only`: trains only parameters whose names contain + `action_expert`. It requires `policy.action_mode=continuous`. +- `policy.enable_lora_vlm`: enables LoRA on VLM linear layers. Use + `policy.enable_lora_action_expert=true` only if LoRA should also cover action + expert linear layers. When `policy.enable_lora_action_expert=false`, the + action expert base weights remain fully trainable while the VLM is trained + through LoRA adapters. When `policy.enable_lora_action_expert=true`, the + action expert is also adapter-tuned instead of fully fine-tuned. +- `policy.enable_knowledge_insulation`: when `true`, detaches action-expert + context K/V states before the action loss. The default is `false`. +- `policy.chunk_size`: action horizon used by the policy. For LIBERO we use + `10`. This LeRobot port overrides the loaded checkpoint's + `max_action_horizon` with this value. +- `policy.n_action_steps`: number of actions consumed from each predicted + chunk before querying the policy again. For LIBERO, set it to `chunk_size`. +- `policy.setup_type`: text inserted into the prompt to describe the robot and + scene, e.g. `single franka robotic arm in libero`. More examples are listed + in the `metadata_by_tag` entries of + [`norm_stats.json`](https://huggingface.co/allenai/MolmoAct2/blob/main/norm_stats.json). +- `policy.control_mode`: text inserted into the prompt to describe the action + space, e.g. `delta end-effector pose` or `absolute joint pose`. +- `policy.image_keys`: ordered LeRobot image observation keys passed to the + processor. +- `policy.model_dtype`: checkpoint/forward dtype, one of `float32`, + `bfloat16`, or `float16`. Use `bfloat16` for normal training. +- `policy.num_flow_timesteps`: number of flow-matching timesteps sampled per + example during training. We use `8` for fine-tuning. +- `policy.num_inference_steps`: optional override for continuous action + generation steps at inference time. +- `policy.gradient_checkpointing`: enables checkpointing in the VLM/action path + to reduce activation memory. +- `policy.freeze_embedding`: freezes input embeddings. The default is `true`. +- `policy.normalize_gripper`: controls whether gripper dimensions are included + in state/action quantile normalization. The default is `false`. +- `policy.normalize_language`: normalizes task strings before prompt + construction. The default is `true`. +- `policy.mask_action_dim_padding`: masks padded dimensions in the flow loss. + Released checkpoints use `policy.expected_max_action_dim=32`. +- `policy.max_sequence_length`: optional manual sequence cap. Leave unset to + infer it from images, state dimension, action dimension, action horizon, and + discrete-action mode. + +### Learning Rates + +MolmoAct2 uses parameter-group learning rates to match the original MolmoAct2 +fine-tuning experiments. + +- Full fine-tuning uses `policy.optimizer_lr=1e-5` for the VLM, + `policy.optimizer_vit_lr=5e-6` for the vision tower, + `policy.optimizer_connector_lr=5e-6` for image connector layers, and + `policy.optimizer_action_expert_lr=5e-5` for the action expert. +- LoRA VLM fine-tuning sets the VLM, vision, and connector LoRA parameter + groups to `5e-5` when `policy.enable_lora_vlm=true`. By default, + `policy.enable_lora_action_expert=false`, so the action expert is still fully + fine-tuned with `policy.optimizer_action_expert_lr`. If + `policy.enable_lora_action_expert=true`, the action expert is trained through + LoRA adapters instead. +- Action-expert-only fine-tuning trains only the action expert and uses + `policy.optimizer_action_expert_lr=5e-5`. + +You can override the full fine-tuning and action-expert learning rates with +`policy.optimizer_lr`, `policy.optimizer_vit_lr`, +`policy.optimizer_connector_lr`, and `policy.optimizer_action_expert_lr`. +Scheduler settings can be changed with `policy.scheduler_warmup_steps`, +`policy.scheduler_decay_steps`, and `policy.scheduler_decay_lr`. + +### Dataset Quantile Statistics + +MolmoAct2 defaults to quantile normalization for state and action features. If +your dataset has not been converted with quantile statistics, you can add them +with: + +```bash +python src/lerobot/datasets/v30/augment_dataset_quantile_stats.py \ + --repo-id=your_dataset +``` + +Alternatively, train MolmoAct2 with mean/std normalization: + +```bash +--policy.normalization_mapping='{"ACTION": "MEAN_STD", "STATE": "MEAN_STD", "VISUAL": "IDENTITY"}' +``` + +## Evaluation + +Evaluation also supports both LeRobot-saved checkpoints and original MolmoAct2 +HF checkpoints. For LIBERO replication, keep the EGL rendering environment +fixed and use `policy.per_episode_seed=true`. + +**Important:** We found that `num_steps_wait=10` does not reliably let the +LIBERO scene stabilize and can degrade measured success. All LIBERO evaluation +results reported here use `num_steps_wait=50`. + +### Evaluation With LeRobot MolmoAct2 Weight + +Use `policy.path` for a checkpoint saved by LeRobot. The saved processor and +normalization statistics are restored together with the model. + +```bash +export MUJOCO_GL=egl +export PYOPENGL_PLATFORM=egl +export OMP_NUM_THREADS=1 +export MKL_NUM_THREADS=1 + +lerobot-eval \ + --policy.path=allenai/MolmoAct2-LIBERO-LeRobot \ + --policy.inference_action_mode=continuous \ + --policy.model_dtype=bfloat16 \ + --policy.use_amp=true \ + --policy.enable_inference_cuda_graph=true \ + --policy.device=cuda \ + --policy.per_episode_seed=true \ + --policy.eval_seed=1000 \ + --env.type=libero \ + --env.task=libero_10,libero_goal,libero_object,libero_spatial \ + --env.camera_name_mapping='{"agentview_image":"image","robot0_eye_in_hand_image":"wrist_image"}' \ + --eval.batch_size=1 \ + --eval.n_episodes=50 \ + --seed=1000 +``` + +### Evaluation With Original MolmoAct2 Weight + +You can evaluate a released Hugging Face checkpoint directly without first +converting it to a LeRobot checkpoint. In this case, set +`policy.checkpoint_path` to the HF model repo and provide `policy.norm_tag`. +For LIBERO, `policy.norm_tag=libero` loads the LIBERO action/state +normalization statistics, action horizon, prompt metadata, and image-key order +from the checkpoint's `norm_stats.json`. + +To fully replicate the MolmoAct2 paper results with released Hugging Face +checkpoints, we recommend using the v0.5.1-pinned +[`allenai/lerobot` `molmoact2-hf-inference`](https://github.com/allenai/lerobot/tree/molmoact2-hf-inference) +branch. That branch matches the original evaluation settings used for the +reported numbers. + +```bash +export MUJOCO_GL=egl +export PYOPENGL_PLATFORM=egl +export OMP_NUM_THREADS=1 +export MKL_NUM_THREADS=1 + +lerobot-eval \ + --policy.type=molmoact2 \ + --policy.checkpoint_path=allenai/MolmoAct2-LIBERO \ + --policy.norm_tag=libero \ + --policy.inference_action_mode=continuous \ + --policy.model_dtype=float32 \ + --policy.use_amp=false \ + --policy.enable_inference_cuda_graph=true \ + --policy.device=cuda \ + --policy.per_episode_seed=true \ + --policy.eval_seed=1000 \ + --env.type=libero \ + --env.task=libero_goal \ + --env.camera_name_mapping='{"agentview_image":"image","robot0_eye_in_hand_image":"wrist_image"}' \ + --eval.batch_size=1 \ + --eval.n_episodes=50 \ + --seed=1000 +``` + +Use `--env.task=libero_10,libero_goal,libero_object,libero_spatial` to run the +full LIBERO suite. The same command works for other released MolmoAct2 +checkpoints as long as the requested `policy.norm_tag` exists in that +checkpoint's `norm_stats.json`. + +### Common Evaluation Options + +- `policy.inference_action_mode`: required for rollout. Use `continuous` for + flow-matching inference or `discrete` for action-token inference. It must be + compatible with the training-time `policy.action_mode` saved in the + checkpoint. +- `policy.path`: LeRobot checkpoint path or Hub repo. Use this for checkpoints + saved by LeRobot. +- `policy.checkpoint_path`: original MolmoAct2 HF checkpoint path or Hub repo. + Use this with `policy.type=molmoact2` and `policy.norm_tag`. +- `policy.norm_tag`: selects normalization statistics, prompt metadata, + image-key order, and action horizon from the original checkpoint's + `norm_stats.json`. It is required for direct original-HF checkpoint + evaluation. +- `policy.model_dtype`: model load/forward dtype. Use `bfloat16` for normal + GPU evaluation. Use `float32` only when you explicitly want fp32 inference. +- `policy.use_amp`: runs the policy forward under autocast during eval. For + `model_dtype=bfloat16`, keep this enabled. +- `policy.enable_inference_cuda_graph`: enables the MolmoAct2 inference CUDA + graph path for faster repeated continuous-action rollout. +- `policy.per_episode_seed` and `policy.eval_seed`: make stochastic continuous + action generation deterministic per episode for replication. +- `env.task`: comma-separated LIBERO suites or a single suite. Use + `libero_10,libero_goal,libero_object,libero_spatial` for the full benchmark. +- `env.camera_name_mapping`: maps LIBERO camera names to the image keys expected + by the policy processor. + +## Performance Results + +### LIBERO Benchmark Results + +MolmoAct2 has demonstrated strong performance on the LIBERO benchmark suite. To +compare and test its LeRobot implementation, we fine-tuned +[`allenai/MolmoAct2-LIBERO`](https://huggingface.co/allenai/MolmoAct2-LIBERO) +for an additional 10k steps on the LIBERO dataset with per-GPU batch size 32 on +8 H100 GPUs, then compared the results to the original MolmoAct2 reference +results. + +The LeRobot fine-tuned checkpoint reported here is available at +[`allenai/MolmoAct2-LIBERO-LeRobot`](https://huggingface.co/allenai/MolmoAct2-LIBERO-LeRobot) +and was trained on +[`allenai/MolmoAct2-LIBERO-Dataset`](https://huggingface.co/datasets/allenai/MolmoAct2-LIBERO-Dataset). + +| Benchmark | LeRobot Implementation | MolmoAct2 Original | +| -------------- | ---------------------: | -----------------: | +| LIBERO Spatial | 98.4% | 97.8% | +| LIBERO Object | 100.0% | 100.0% | +| LIBERO Goal | 98.0% | 97.8% | +| LIBERO 10 | 96.6% | 93.2% | +| Average | 98.25% | 97.20% | + +These results demonstrate MolmoAct2's strong performance across diverse robotic +manipulation tasks. To reproduce them, follow the instructions in the LIBERO +evaluation section. + +## Differences From the Original Implementation + +This LeRobot port is intended to match MolmoAct2 behavior while using LeRobot's +dataset, training, evaluation, checkpoint, and logging infrastructure. The main +differences from the original training repository are: + +- The original paper training stack loads the model in fp32 and trains under + mixed precision. This LeRobot port usually loads the checkpoint directly in + `policy.model_dtype=bfloat16` for lower memory use. +- The original repository uses its own FSDP/model-parallel training path. The + LeRobot port uses the standard LeRobot/Accelerate training path and has not + been tested for multi-node training. +- The original repository supports sequence packing. The LeRobot port trains on + one LeRobot sample per item and pads to an inferred fixed sequence budget. +- The LeRobot port follows LeRobot's optimizer, scheduler, checkpoint saving, + dataset transforms, image augmentation, and Weights & Biases logging + conventions. +- The original training path supports mixed action horizons by padding to + `max_action_horizon` and masking padded horizon slots in the action expert + self-attention. This is useful when training across datasets with different + control frequencies. The LeRobot port currently targets single-dataset + fine-tuning, so `policy.chunk_size` overrides the checkpoint + `max_action_horizon` and horizon masking is not implemented yet. Support for + this mixed-horizon path is planned. + +## Citation + +```bibtex +@misc{fang2026molmoact2actionreasoningmodels, + title={MolmoAct2: Action Reasoning Models for Real-world Deployment}, + author={Haoquan Fang and Jiafei Duan and Donovan Clay and Sam Wang and Shuo Liu and Weikai Huang and Xiang Fan and Wei-Chuan Tsai and Shirui Chen and Yi Ru Wang and Shanli Xing and Jaemin Cho and Jae Sung Park and Ainaz Eftekhar and Peter Sushko and Karen Farley and Angad Wadhwa and Cole Harrison and Winson Han and Ying-Chun Lee and Eli VanderBilt and Rose Hendrix and Suveen Ellawela and Lucas Ngoo and Joyce Chai and Zhongzheng Ren and Ali Farhadi and Dieter Fox and Ranjay Krishna}, + year={2026}, + eprint={2605.02881}, + archivePrefix={arXiv}, + primaryClass={cs.RO}, + url={https://arxiv.org/abs/2605.02881}, +} +``` + +## License + +This model is licensed under Apache 2.0. It is intended for research and +educational use in accordance with +[Ai2's Responsible Use Guidelines](https://allenai.org/responsible-use), +consistent with [allenai/molmoact2](https://github.com/allenai/molmoact2). diff --git a/docs/source/policy_molmoact2_README.md b/docs/source/policy_molmoact2_README.md new file mode 100644 index 000000000..df3a6341e --- /dev/null +++ b/docs/source/policy_molmoact2_README.md @@ -0,0 +1,39 @@ +# MolmoAct2 + +This repository contains the LeRobot policy implementation of +[MolmoAct2](https://allenai.org/blog/molmoact2), ported into LeRobot for +training, evaluation, checkpointing, and dataset compatibility. + +This implementation currently supports training and evaluation for the regular +MolmoAct2 model. MolmoAct2-Think, which supports adaptive depth reasoning, is +not included in this LeRobot policy yet and is coming soon. + +For the original MolmoAct2 training code used for the experiments reported in +the paper, see [allenai/molmoact2](https://github.com/allenai/molmoact2). + +## LIBERO Evaluation + +Important: we found that `num_steps_wait=10` does not reliably let the LIBERO +scene stabilize and can degrade measured success. All LIBERO evaluation results +reported for this LeRobot implementation use `num_steps_wait=50`. + +## Citation + +```bibtex +@misc{fang2026molmoact2actionreasoningmodels, + title={MolmoAct2: Action Reasoning Models for Real-world Deployment}, + author={Haoquan Fang and Jiafei Duan and Donovan Clay and Sam Wang and Shuo Liu and Weikai Huang and Xiang Fan and Wei-Chuan Tsai and Shirui Chen and Yi Ru Wang and Shanli Xing and Jaemin Cho and Jae Sung Park and Ainaz Eftekhar and Peter Sushko and Karen Farley and Angad Wadhwa and Cole Harrison and Winson Han and Ying-Chun Lee and Eli VanderBilt and Rose Hendrix and Suveen Ellawela and Lucas Ngoo and Joyce Chai and Zhongzheng Ren and Ali Farhadi and Dieter Fox and Ranjay Krishna}, + year={2026}, + eprint={2605.02881}, + archivePrefix={arXiv}, + primaryClass={cs.RO}, + url={https://arxiv.org/abs/2605.02881}, +} +``` + +## License + +This model is licensed under Apache 2.0. It is intended for research and +educational use in accordance with +[Ai2's Responsible Use Guidelines](https://allenai.org/responsible-use), +consistent with [allenai/molmoact2](https://github.com/allenai/molmoact2). diff --git a/docs/source/robometer.mdx b/docs/source/robometer.mdx new file mode 100644 index 000000000..5af6882d3 --- /dev/null +++ b/docs/source/robometer.mdx @@ -0,0 +1,185 @@ +# ROBOMETER + +ROBOMETER is a **general-purpose video-language robotic reward model**. It predicts dense, frame-level task progress and frame-level success from a trajectory video and a task description. + +**Paper**: [ROBOMETER: Scaling General-Purpose Robotic Reward Models via Trajectory Comparisons](https://arxiv.org/abs/2603.02115) +**Project**: [robometer.github.io](https://robometer.github.io/) +**Original code**: [github.com/robometer/robometer](https://github.com/robometer/robometer) +**Checkpoint**: [lerobot/Robometer-4B](https://huggingface.co/lerobot/Robometer-4B) + +## Overview + +ROBOMETER builds on `Qwen/Qwen3-VL-4B-Instruct` and adds three lightweight prediction heads: + +- **Progress head**: predicts per-frame task progress in `[0, 1]`. +- **Success head**: predicts per-frame task success probability. +- **Preference head**: predicts which of two trajectories better completes the task during training. + +The paper trains ROBOMETER with a composite objective: + +```text +L = L_pref + L_prog + L_succ +``` + +The LeRobot integration is currently **inference-only**. It preserves the preference head so that the published `Robometer-4B` checkpoint loads without remapping, but `compute_reward()` queries the progress or success head only. + +## What the LeRobot Integration Covers + +- Standard `reward_model.type=robometer` configuration through LeRobot. +- Qwen3-VL image and text preprocessing through `RobometerEncoderProcessorStep`. +- LeRobot reward-model save/load APIs through `PreTrainedRewardModel`. +- Dense, frame-level progress and success predictions internally. +- A scalar reward through `compute_reward()` for downstream LeRobot reward-model usage. + +This page focuses on using the published ROBOMETER checkpoint as a zero-shot reward model. Training ROBOMETER from scratch is outside the current LeRobot integration. + +## Installation Requirements + +1. Install LeRobot by following the [Installation Guide](./installation). +2. Install the ROBOMETER dependencies: + +```bash +pip install -e ".[robometer]" +``` + +If you use `uv` directly from a source checkout: + +```bash +uv sync --extra robometer +``` + +ROBOMETER uses a Qwen3-VL-4B backbone, so GPU inference is strongly recommended. + +## Model Inputs and Outputs + +ROBOMETER expects: + +- A trajectory video or sequence of frames. +- A natural-language task description. + +In LeRobot datasets, the preprocessor reads: + +| Config field | Default | Meaning | +| ------------------------- | ------------------------ | ----------------------------------------------------- | +| `reward_model.image_key` | `observation.images.top` | Camera/video observation used by ROBOMETER | +| `reward_model.task_key` | `task` | Key in complementary data that stores the task string | +| `reward_model.max_frames` | `8` | Maximum number of frames passed to ROBOMETER | + +The model predicts per-frame progress and success internally. The LeRobot reward API returns a scalar per sample: + +- `reward_output="progress"` (default): return the last-frame progress, clamped to `[0, 1]`. +- `reward_output="success"`: return `1.0` if the last-frame success probability is above `success_threshold`, otherwise `0.0`. + +## Usage + +### Load the Reward Model Directly + +```python +from lerobot.rewards.robometer import RobometerConfig, RobometerRewardModel + +cfg = RobometerConfig( + pretrained_path="lerobot/Robometer-4B", + device="cuda", + reward_output="progress", +) +reward_model = RobometerRewardModel.from_pretrained(cfg.pretrained_path, config=cfg) +``` + +### Encode Frames and Compute a Reward + +For a direct Python call, provide frames as `uint8` arrays with shape `(T, H, W, C)` and a task string: + +```python +from lerobot.rewards.robometer.modeling_robometer import ROBOMETER_FEATURE_PREFIX +from lerobot.rewards.robometer.processor_robometer import RobometerEncoderProcessorStep + +# frames: np.ndarray, shape (T, H, W, C), dtype uint8 +# task: str +encoder = RobometerEncoderProcessorStep( + base_model_id=cfg.base_model_id, + use_multi_image=cfg.use_multi_image, + use_per_frame_progress_token=cfg.use_per_frame_progress_token, + max_frames=cfg.max_frames, +) + +encoded = encoder.encode_samples([(frames, task)]) +batch = {f"{ROBOMETER_FEATURE_PREFIX}{key}": value for key, value in encoded.items()} + +reward = reward_model.compute_reward(batch) +``` + +`reward` is a tensor of shape `(batch_size,)`. + +### Use the Reward Factory + +You can also instantiate ROBOMETER through the reward factory: + +```python +from lerobot.rewards import make_reward_model, make_reward_model_config, make_reward_pre_post_processors + +cfg = make_reward_model_config( + "robometer", + pretrained_path="lerobot/Robometer-4B", + device="cuda", + image_key="observation.images.top", +) +reward_model = make_reward_model(cfg) +preprocessor, postprocessor = make_reward_pre_post_processors(cfg) +``` + +The preprocessor writes Qwen-VL tensors under the `observation.robometer.*` namespace, and `compute_reward()` reads those encoded tensors. + +## Configuration Notes + +### Backbone and Vocabulary + +The published checkpoint uses a Qwen3-VL-4B backbone. ROBOMETER adds five special tokens to the tokenizer in a fixed order: + +```text +<|split_token|> +<|reward_token|> +<|pref_token|> +<|sim_token|> +<|prog_token|> +``` + +`<|prog_token|>` is inserted after each frame and is the hidden-state position used for per-frame progress and success prediction. `<|split_token|>` and `<|pref_token|>` are used by the paper's pairwise trajectory preference objective. `<|reward_token|>` and `<|sim_token|>` are preserved for checkpoint compatibility. + +The LeRobot config stores a serialized `vlm_config` with the post-resize vocabulary so the model can reload from `config.json` without downloading the base Qwen weights first. For `Qwen/Qwen3-VL-4B-Instruct`, the tokenizer length is `151669`, and the five ROBOMETER tokens produce the checkpoint vocabulary size `151674`. + +### Progress Prediction + +In the published checkpoint, progress is discrete. The progress head outputs logits over `progress_discrete_bins=10` uniformly spaced bin centers in `[0, 1]`. LeRobot converts these logits into a continuous value by applying a softmax and taking the expectation over bin centers, matching the upstream ROBOMETER implementation. + +### Success Prediction + +The success head outputs raw logits per frame. LeRobot converts them to probabilities with `sigmoid`. When `reward_output="success"`, `compute_reward()` thresholds the last-frame success probability using `success_threshold`. + +## Limitations + +- The current LeRobot integration is inference-only; it does not implement ROBOMETER training or preference-pair training. +- `compute_reward()` returns a scalar per sample for the LeRobot reward-model API, even though ROBOMETER predicts per-frame progress and success internally. +- ROBOMETER is video-language based; it does not use privileged robot state such as contact forces or object poses. + +## References + +- [ROBOMETER project](https://robometer.github.io/) +- [ROBOMETER paper](https://arxiv.org/abs/2603.02115) +- [Original ROBOMETER code](https://github.com/robometer/robometer) +- [Published ROBOMETER-4B checkpoint](https://huggingface.co/lerobot/Robometer-4B) +- [Qwen3-VL-4B-Instruct](https://huggingface.co/Qwen/Qwen3-VL-4B-Instruct) + +## Citation + +```bibtex +@inproceedings{liang2026robometer, +title = {Robometer: Scaling General-Purpose Robotic Reward Models via Trajectory Comparisons}, +author={Anthony Liang and Yigit Korkmaz and Jiahui Zhang and Minyoung Hwang and Abrar Anwar and Sidhant Kaushik and Aditya Shah and Alex S. Huang and Luke Zettlemoyer and Dieter Fox and Yu Xiang and Anqi Li and Andreea Bobu and Abhishek Gupta and Stephen Tu and Erdem Biyik and Jesse Zhang}, +year={2026}, +booktitle={Robotics: Science and Systems 2026}, +} +``` + +## License + +This LeRobot integration follows the **Apache 2.0 License** used by LeRobot. Check the upstream ROBOMETER code and model pages for the licenses of the original implementation and released checkpoints. diff --git a/docs/source/topreward.mdx b/docs/source/topreward.mdx new file mode 100644 index 000000000..f84fbed49 --- /dev/null +++ b/docs/source/topreward.mdx @@ -0,0 +1,177 @@ +# TOPReward + +TOPReward is a **zero-shot reward model** that extracts token log-probabilities from an off-the-shelf vision-language model (VLM) as a robotic reward signal. Given a video trajectory and a task instruction, it returns the VLM's log-likelihood that the instruction is true — no fine-tuning required. + +**Paper**: [TOPReward: Token Probabilities as Hidden Zero-Shot Rewards for Robotics](https://arxiv.org/abs/2602.19313) +**Project**: [topreward.github.io](https://topreward.github.io/webpage/) +**Original code**: [github.com/TOPReward/TOPReward](https://github.com/TOPReward/TOPReward) +**Default backbone**: [Qwen/Qwen3-VL-8B-Instruct](https://huggingface.co/Qwen/Qwen3-VL-8B-Instruct) + +## Overview + +TOPReward asks a generic VLM how likely a task instruction is, **conditioned on the video** of a robot trying to complete that task. Concretely, given: + +- A trajectory video (a sequence of frames). +- A task instruction (e.g. _"open the drawer"_). + +it builds a chat prompt of the form + +```text +