mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-08 10:32:00 +00:00
add molmoact2 policy
This commit is contained in:
hq-fang
@@ -59,6 +59,8 @@
|
||||
title: π₀-FAST (Pi0Fast)
|
||||
- local: pi05
|
||||
title: π₀.₅ (Pi05)
|
||||
- local: molmoact2
|
||||
title: MolmoAct2
|
||||
- local: eo1
|
||||
title: EO-1
|
||||
- local: groot
|
||||
|
||||
@@ -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_entity> \
|
||||
--wandb.project=<wandb_project> \
|
||||
--job_name=<job_name> \
|
||||
--output_dir=outputs/<job_name> \
|
||||
--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_entity> \
|
||||
--wandb.project=<wandb_project> \
|
||||
--job_name=<job_name> \
|
||||
--output_dir=outputs/<job_name> \
|
||||
--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).
|
||||
@@ -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).
|
||||
+4
-1
@@ -198,6 +198,7 @@ wallx = [
|
||||
"lerobot[qwen-vl-utils-dep]",
|
||||
]
|
||||
pi = ["lerobot[transformers-dep]", "lerobot[scipy-dep]"]
|
||||
molmoact2 = ["lerobot[transformers-dep]", "lerobot[peft-dep]"]
|
||||
smolvla = ["lerobot[transformers-dep]", "num2words>=0.5.14,<0.6.0", "accelerate>=1.7.0,<2.0.0"]
|
||||
multi_task_dit = ["lerobot[transformers-dep]", "lerobot[diffusers-dep]"]
|
||||
groot = [
|
||||
@@ -274,6 +275,7 @@ all = [
|
||||
"lerobot[multi_task_dit]",
|
||||
"lerobot[wallx]",
|
||||
"lerobot[pi]",
|
||||
"lerobot[molmoact2]",
|
||||
"lerobot[smolvla]",
|
||||
# "lerobot[groot]", TODO(Steven): Gr00t requires specific installation instructions for flash-attn
|
||||
"lerobot[xvla]",
|
||||
@@ -404,7 +406,8 @@ default.extend-ignore-identifiers-re = [
|
||||
"thw",
|
||||
"inpt",
|
||||
"ROBOTIS",
|
||||
"OT_VALUE"
|
||||
"OT_VALUE",
|
||||
"VanderBilt"
|
||||
]
|
||||
|
||||
# TODO: Uncomment when ready to use
|
||||
|
||||
@@ -20,6 +20,7 @@ from .eo1.configuration_eo1 import EO1Config as EO1Config
|
||||
from .factory import get_policy_class, make_policy, make_policy_config, make_pre_post_processors
|
||||
from .gaussian_actor.configuration_gaussian_actor import GaussianActorConfig as GaussianActorConfig
|
||||
from .groot.configuration_groot import GrootConfig as GrootConfig
|
||||
from .molmoact2.configuration_molmoact2 import MolmoAct2Config as MolmoAct2Config
|
||||
from .multi_task_dit.configuration_multi_task_dit import MultiTaskDiTConfig as MultiTaskDiTConfig
|
||||
from .pi0.configuration_pi0 import PI0Config as PI0Config
|
||||
from .pi0_fast.configuration_pi0_fast import PI0FastConfig as PI0FastConfig
|
||||
@@ -43,6 +44,7 @@ __all__ = [
|
||||
"EO1Config",
|
||||
"GaussianActorConfig",
|
||||
"GrootConfig",
|
||||
"MolmoAct2Config",
|
||||
"MultiTaskDiTConfig",
|
||||
"PI0Config",
|
||||
"PI0FastConfig",
|
||||
|
||||
@@ -49,6 +49,7 @@ from .diffusion.configuration_diffusion import DiffusionConfig
|
||||
from .eo1.configuration_eo1 import EO1Config
|
||||
from .gaussian_actor.configuration_gaussian_actor import GaussianActorConfig
|
||||
from .groot.configuration_groot import GrootConfig
|
||||
from .molmoact2.configuration_molmoact2 import MolmoAct2Config
|
||||
from .multi_task_dit.configuration_multi_task_dit import MultiTaskDiTConfig
|
||||
from .pi0.configuration_pi0 import PI0Config
|
||||
from .pi05.configuration_pi05 import PI05Config
|
||||
@@ -88,7 +89,8 @@ def get_policy_class(name: str) -> type[PreTrainedPolicy]:
|
||||
|
||||
Args:
|
||||
name: The name of the policy. Supported names are "tdmpc", "diffusion", "act",
|
||||
"multi_task_dit", "vqbet", "pi0", "pi05", "gaussian_actor", "smolvla", "wall_x".
|
||||
"multi_task_dit", "vqbet", "pi0", "pi05", "gaussian_actor", "smolvla", "wall_x",
|
||||
"molmoact2".
|
||||
Returns:
|
||||
The policy class corresponding to the given name.
|
||||
|
||||
@@ -151,6 +153,10 @@ def get_policy_class(name: str) -> type[PreTrainedPolicy]:
|
||||
from .eo1.modeling_eo1 import EO1Policy
|
||||
|
||||
return EO1Policy
|
||||
elif name == "molmoact2":
|
||||
from .molmoact2.modeling_molmoact2 import MolmoAct2Policy
|
||||
|
||||
return MolmoAct2Policy
|
||||
else:
|
||||
try:
|
||||
return _get_policy_cls_from_policy_name(name=name)
|
||||
@@ -168,7 +174,7 @@ def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
|
||||
Args:
|
||||
policy_type: The type of the policy. Supported types include "tdmpc",
|
||||
"multi_task_dit", "diffusion", "act", "vqbet", "pi0", "pi05", "gaussian_actor",
|
||||
"smolvla", "wall_x".
|
||||
"smolvla", "wall_x", "molmoact2".
|
||||
**kwargs: Keyword arguments to be passed to the configuration class constructor.
|
||||
|
||||
Returns:
|
||||
@@ -203,6 +209,8 @@ def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
|
||||
return WallXConfig(**kwargs)
|
||||
elif policy_type == "eo1":
|
||||
return EO1Config(**kwargs)
|
||||
elif policy_type == "molmoact2":
|
||||
return MolmoAct2Config(**kwargs)
|
||||
else:
|
||||
try:
|
||||
config_cls = PreTrainedConfig.get_choice_class(policy_type)
|
||||
@@ -231,6 +239,7 @@ class ProcessorConfigKwargs(TypedDict, total=False):
|
||||
preprocessor_overrides: dict[str, Any] | None
|
||||
postprocessor_overrides: dict[str, Any] | None
|
||||
dataset_stats: dict[str, dict[str, torch.Tensor]] | None
|
||||
dataset_meta: Any | None
|
||||
|
||||
|
||||
def make_pre_post_processors(
|
||||
@@ -285,6 +294,25 @@ def make_pre_post_processors(
|
||||
kwargs["preprocessor_overrides"] = preprocessor_overrides
|
||||
kwargs["postprocessor_overrides"] = postprocessor_overrides
|
||||
|
||||
if isinstance(policy_cfg, MolmoAct2Config):
|
||||
from .molmoact2 import processor_molmoact2 # noqa: F401
|
||||
|
||||
preprocessor_overrides = dict(kwargs.get("preprocessor_overrides", {}))
|
||||
if "normalizer_processor" in preprocessor_overrides:
|
||||
preprocessor_overrides.setdefault(
|
||||
"molmoact2_masked_normalizer",
|
||||
preprocessor_overrides.pop("normalizer_processor"),
|
||||
)
|
||||
kwargs["preprocessor_overrides"] = preprocessor_overrides
|
||||
|
||||
postprocessor_overrides = dict(kwargs.get("postprocessor_overrides", {}))
|
||||
if "unnormalizer_processor" in postprocessor_overrides:
|
||||
postprocessor_overrides.setdefault(
|
||||
"molmoact2_masked_unnormalizer",
|
||||
postprocessor_overrides.pop("unnormalizer_processor"),
|
||||
)
|
||||
kwargs["postprocessor_overrides"] = postprocessor_overrides
|
||||
|
||||
preprocessor = PolicyProcessorPipeline.from_pretrained(
|
||||
pretrained_model_name_or_path=pretrained_path,
|
||||
config_filename=kwargs.get(
|
||||
@@ -414,6 +442,15 @@ def make_pre_post_processors(
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
)
|
||||
|
||||
elif isinstance(policy_cfg, MolmoAct2Config):
|
||||
from .molmoact2.processor_molmoact2 import make_molmoact2_pre_post_processors
|
||||
|
||||
processors = make_molmoact2_pre_post_processors(
|
||||
config=policy_cfg,
|
||||
dataset_stats=kwargs.get("dataset_stats"),
|
||||
dataset_meta=kwargs.get("dataset_meta"),
|
||||
)
|
||||
|
||||
else:
|
||||
try:
|
||||
processors = _make_processors_from_policy_config(
|
||||
@@ -499,6 +536,10 @@ def make_policy(
|
||||
action_names = ds_meta.features.get(ACTION, {}).get("names")
|
||||
if action_names is not None:
|
||||
cfg.action_feature_names = list(action_names)
|
||||
if ds_meta is not None:
|
||||
set_dataset_feature_metadata = getattr(cfg, "set_dataset_feature_metadata", None)
|
||||
if callable(set_dataset_feature_metadata):
|
||||
set_dataset_feature_metadata(ds_meta.features)
|
||||
|
||||
kwargs["config"] = cfg
|
||||
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
../../../../docs/source/policy_molmoact2_README.md
|
||||
@@ -0,0 +1,15 @@
|
||||
from .configuration_molmoact2 import MolmoAct2Config
|
||||
|
||||
__all__ = ["MolmoAct2Config", "MolmoAct2Policy", "make_molmoact2_pre_post_processors"]
|
||||
|
||||
|
||||
def __getattr__(name):
|
||||
if name == "MolmoAct2Policy":
|
||||
from .modeling_molmoact2 import MolmoAct2Policy
|
||||
|
||||
return MolmoAct2Policy
|
||||
if name == "make_molmoact2_pre_post_processors":
|
||||
from .processor_molmoact2 import make_molmoact2_pre_post_processors
|
||||
|
||||
return make_molmoact2_pre_post_processors
|
||||
raise AttributeError(name)
|
||||
@@ -0,0 +1,324 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
from lerobot.configs import FeatureType, NormalizationMode, PolicyFeature, PreTrainedConfig
|
||||
from lerobot.optim import (
|
||||
AdamWConfig,
|
||||
CosineDecayWithWarmupSchedulerConfig,
|
||||
LRSchedulerConfig,
|
||||
OptimizerConfig,
|
||||
)
|
||||
from lerobot.utils.constants import ACTION, OBS_STATE
|
||||
|
||||
from ..rtc.configuration_rtc import RTCConfig
|
||||
|
||||
MOLMOACT2_DEFAULT_NUM_IMAGES = 2
|
||||
MOLMOACT2_IMAGE_TOKENS_PER_IMAGE = 196
|
||||
MOLMOACT2_FIXED_PROMPT_TOKEN_BUDGET = 80
|
||||
MOLMOACT2_TASK_TOKEN_BUDGET = 32
|
||||
MOLMOACT2_SEQUENCE_LENGTH_MARGIN = 32
|
||||
MOLMOACT2_SEQUENCE_LENGTH_MULTIPLE = 64
|
||||
MOLMOACT2_DISCRETE_ACTION_WRAPPER_TOKENS = 4
|
||||
MOLMOACT2_MIN_DISCRETE_ACTION_TOKENS_PER_STEP = 6
|
||||
MOLMOACT2_DISCRETE_ACTION_TOKENS_PER_DIM = 0.95
|
||||
|
||||
|
||||
@LRSchedulerConfig.register_subclass("molmoact2_cosine_decay_with_warmup")
|
||||
@dataclass
|
||||
class MolmoAct2CosineDecayWithWarmupSchedulerConfig(CosineDecayWithWarmupSchedulerConfig):
|
||||
"""MolmoAct2-local cosine scheduler with optional decay-step auto-match.
|
||||
|
||||
LeRobot's generic cosine scheduler keeps an explicit integer decay length.
|
||||
For MolmoAct2, leaving num_decay_steps unset means "decay across this run's
|
||||
training steps"; build() is the first point where num_training_steps is known.
|
||||
"""
|
||||
|
||||
num_decay_steps: int | None
|
||||
|
||||
def build(self, optimizer, num_training_steps: int):
|
||||
return CosineDecayWithWarmupSchedulerConfig(
|
||||
peak_lr=self.peak_lr,
|
||||
decay_lr=self.decay_lr,
|
||||
num_warmup_steps=self.num_warmup_steps,
|
||||
num_decay_steps=num_training_steps if self.num_decay_steps is None else self.num_decay_steps,
|
||||
).build(optimizer, num_training_steps=num_training_steps)
|
||||
|
||||
|
||||
def _round_up(value: int, multiple: int) -> int:
|
||||
return int(math.ceil(value / multiple) * multiple)
|
||||
|
||||
|
||||
def infer_molmoact2_max_sequence_length(
|
||||
*,
|
||||
num_images: int,
|
||||
state_dim: int,
|
||||
action_dim: int,
|
||||
action_horizon: int,
|
||||
include_discrete_action: bool,
|
||||
) -> int:
|
||||
"""Infer the padded text/image sequence cap from MolmoAct2's fixed token layout."""
|
||||
if num_images < 1:
|
||||
num_images = MOLMOACT2_DEFAULT_NUM_IMAGES
|
||||
if state_dim < 0:
|
||||
state_dim = 0
|
||||
if action_dim < 1:
|
||||
action_dim = 1
|
||||
if action_horizon < 1:
|
||||
action_horizon = 1
|
||||
|
||||
image_tokens = num_images * MOLMOACT2_IMAGE_TOKENS_PER_IMAGE
|
||||
prompt_tokens = (
|
||||
MOLMOACT2_FIXED_PROMPT_TOKEN_BUDGET
|
||||
+ MOLMOACT2_TASK_TOKEN_BUDGET
|
||||
+ state_dim
|
||||
+ MOLMOACT2_SEQUENCE_LENGTH_MARGIN
|
||||
)
|
||||
action_tokens = 0
|
||||
if include_discrete_action:
|
||||
action_tokens_per_step = max(
|
||||
MOLMOACT2_MIN_DISCRETE_ACTION_TOKENS_PER_STEP,
|
||||
math.ceil(action_dim * MOLMOACT2_DISCRETE_ACTION_TOKENS_PER_DIM),
|
||||
)
|
||||
action_tokens = MOLMOACT2_DISCRETE_ACTION_WRAPPER_TOKENS + action_horizon * action_tokens_per_step
|
||||
|
||||
return _round_up(
|
||||
image_tokens + prompt_tokens + action_tokens,
|
||||
MOLMOACT2_SEQUENCE_LENGTH_MULTIPLE,
|
||||
)
|
||||
|
||||
|
||||
@PreTrainedConfig.register_subclass("molmoact2")
|
||||
@dataclass
|
||||
class MolmoAct2Config(PreTrainedConfig):
|
||||
"""MolmoAct2 policy backed by the converted HF checkpoint implementation."""
|
||||
|
||||
checkpoint_path: str = "allenai/MolmoAct2"
|
||||
checkpoint_revision: str | None = None
|
||||
checkpoint_force_download: bool = False
|
||||
trust_remote_code: bool = True
|
||||
|
||||
n_obs_steps: int = 1
|
||||
chunk_size: int = 30
|
||||
n_action_steps: int = 30
|
||||
|
||||
action_mode: str = "both"
|
||||
inference_action_mode: str | None = None
|
||||
discrete_action_tokenizer: str = "allenai/MolmoAct2-FAST-Tokenizer"
|
||||
discrete_generation_max_steps: int | None = None
|
||||
norm_tag: str | None = None
|
||||
|
||||
setup_type: str = ""
|
||||
control_mode: str = ""
|
||||
image_keys: list[str] = field(default_factory=list)
|
||||
normalize_language: bool = True
|
||||
add_setup_tokens: bool = True
|
||||
add_control_tokens: bool = True
|
||||
normalize_gripper: bool = False
|
||||
num_state_tokens: int = 256
|
||||
# Leave unset for the default MolmoAct2 sequence budget inferred from the fixed
|
||||
# image/prompt/state/action token layout. Override only for unusual long prompts.
|
||||
max_sequence_length: int | None = None
|
||||
|
||||
# Fixed by released MolmoAct2 checkpoints. We validate this at model load.
|
||||
expected_max_action_dim: int = 32
|
||||
|
||||
# Flow-matching training knobs copied from the original MolmoAct2 training path.
|
||||
num_flow_timesteps: int = 8
|
||||
flow_matching_cutoff: float = 1.0
|
||||
flow_matching_time_offset: float = 0.001
|
||||
flow_matching_time_scale: float = 0.999
|
||||
flow_matching_beta_alpha: float = 1.0
|
||||
flow_matching_beta_beta: float = 1.5
|
||||
num_inference_steps: int | None = None
|
||||
mask_action_dim_padding: bool = True
|
||||
enable_inference_cuda_graph: bool = True
|
||||
# MolmoAct2-local eval option. When enabled, stochastic continuous action
|
||||
# generation uses a rollout-local generator derived from eval_seed.
|
||||
per_episode_seed: bool = False
|
||||
eval_seed: int | None = None
|
||||
rtc_config: RTCConfig | None = None
|
||||
|
||||
# Default is full finetuning with gradients from the action expert flowing into the VLM.
|
||||
enable_lora_vlm: bool = False
|
||||
lora_rank: int = 64
|
||||
lora_alpha: int = 16
|
||||
lora_dropout: float = 0.05
|
||||
lora_bias: str = "none"
|
||||
enable_lora_action_expert: bool = False
|
||||
enable_knowledge_insulation: bool = False
|
||||
freeze_embedding: bool = True
|
||||
train_action_expert_only: bool = False
|
||||
gradient_checkpointing: bool = False
|
||||
|
||||
model_dtype: str = "bfloat16"
|
||||
softmax_auxiliary_loss: bool = True
|
||||
softmax_auxiliary_loss_scale: float = 1e-4
|
||||
discrete_loss_token_weighting: str = "root_subsegments_root_tokens"
|
||||
|
||||
optimizer_lr: float = 1e-5
|
||||
optimizer_vit_lr: float = 5e-6
|
||||
optimizer_connector_lr: float = 5e-6
|
||||
optimizer_action_expert_lr: float = 5e-5
|
||||
optimizer_betas: tuple[float, float] = (0.9, 0.95)
|
||||
optimizer_eps: float = 1e-6
|
||||
optimizer_weight_decay: float = 0.0
|
||||
optimizer_grad_clip_norm: float = 1.0
|
||||
|
||||
scheduler_warmup_steps: int = 200
|
||||
scheduler_decay_steps: int | None = None
|
||||
scheduler_decay_lr: float = 1e-6
|
||||
|
||||
normalization_mapping: dict[str, NormalizationMode] = field(
|
||||
default_factory=lambda: {
|
||||
"VISUAL": NormalizationMode.IDENTITY,
|
||||
"STATE": NormalizationMode.QUANTILES,
|
||||
"ACTION": NormalizationMode.QUANTILES,
|
||||
}
|
||||
)
|
||||
|
||||
input_features: dict[str, PolicyFeature] = field(default_factory=dict)
|
||||
output_features: dict[str, PolicyFeature] = field(default_factory=dict)
|
||||
dataset_feature_names: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
super().__post_init__()
|
||||
if self.action_mode not in {"continuous", "discrete", "both"}:
|
||||
raise ValueError(
|
||||
f"Unsupported action_mode={self.action_mode!r}. "
|
||||
"Expected one of {'continuous', 'discrete', 'both'}."
|
||||
)
|
||||
if self.inference_action_mode not in {None, "continuous", "discrete"}:
|
||||
raise ValueError(
|
||||
f"Unsupported inference_action_mode={self.inference_action_mode!r}. "
|
||||
"Expected one of {None, 'continuous', 'discrete'}."
|
||||
)
|
||||
if self.inference_action_mode == "continuous" and self.action_mode == "discrete":
|
||||
raise ValueError("MolmoAct2 action_mode='discrete' cannot run continuous inference.")
|
||||
if self.inference_action_mode == "discrete" and self.action_mode == "continuous":
|
||||
raise ValueError("MolmoAct2 action_mode='continuous' cannot run discrete inference.")
|
||||
if self.train_action_expert_only and self.action_mode != "continuous":
|
||||
raise ValueError("MolmoAct2 train_action_expert_only requires action_mode='continuous'.")
|
||||
if self.train_action_expert_only and self.enable_lora_vlm:
|
||||
raise ValueError("MolmoAct2 train_action_expert_only is incompatible with enable_lora_vlm.")
|
||||
if self.enable_lora_action_expert and not self.enable_lora_vlm:
|
||||
raise ValueError("MolmoAct2 enable_lora_action_expert requires enable_lora_vlm.")
|
||||
if self.chunk_size < 1:
|
||||
raise ValueError(f"chunk_size must be >= 1, got {self.chunk_size}.")
|
||||
if self.n_action_steps < 1:
|
||||
raise ValueError(f"n_action_steps must be >= 1, got {self.n_action_steps}.")
|
||||
if self.n_action_steps > self.chunk_size:
|
||||
raise ValueError(
|
||||
f"n_action_steps ({self.n_action_steps}) cannot exceed chunk_size ({self.chunk_size})."
|
||||
)
|
||||
if self.expected_max_action_dim != 32:
|
||||
raise ValueError("MolmoAct2 released checkpoints use expected_max_action_dim=32.")
|
||||
if self.model_dtype not in {"float32", "bfloat16", "float16"}:
|
||||
raise ValueError(
|
||||
f"Unsupported model_dtype={self.model_dtype!r}. Expected 'float32', 'bfloat16', or 'float16'."
|
||||
)
|
||||
if self.lora_rank < 1:
|
||||
raise ValueError(f"lora_rank must be >= 1, got {self.lora_rank}.")
|
||||
if self.lora_alpha < 1:
|
||||
raise ValueError(f"lora_alpha must be >= 1, got {self.lora_alpha}.")
|
||||
if not 0 <= self.lora_dropout <= 1:
|
||||
raise ValueError(f"lora_dropout must be in [0, 1], got {self.lora_dropout}.")
|
||||
if self.lora_bias not in {"none", "all", "lora_only"}:
|
||||
raise ValueError(
|
||||
f"Unsupported lora_bias={self.lora_bias!r}. Expected one of 'none', 'all', or 'lora_only'."
|
||||
)
|
||||
if self.discrete_loss_token_weighting not in {
|
||||
"none",
|
||||
"token",
|
||||
"root_tokens",
|
||||
"root_subsegments",
|
||||
"root_subsegments_root_tokens",
|
||||
}:
|
||||
raise ValueError(
|
||||
f"Unsupported discrete_loss_token_weighting={self.discrete_loss_token_weighting!r}."
|
||||
)
|
||||
if self.discrete_generation_max_steps is not None and self.discrete_generation_max_steps < 1:
|
||||
raise ValueError(
|
||||
f"discrete_generation_max_steps must be >= 1 or None, got {self.discrete_generation_max_steps}."
|
||||
)
|
||||
if self.max_sequence_length is not None and self.max_sequence_length < 1:
|
||||
raise ValueError(f"max_sequence_length must be >= 1 or None, got {self.max_sequence_length}.")
|
||||
|
||||
def inferred_max_sequence_length(
|
||||
self,
|
||||
*,
|
||||
num_images: int | None = None,
|
||||
state_dim: int | None = None,
|
||||
action_dim: int | None = None,
|
||||
action_horizon: int | None = None,
|
||||
include_discrete_action: bool | None = None,
|
||||
) -> int:
|
||||
if self.max_sequence_length is not None:
|
||||
return int(self.max_sequence_length)
|
||||
|
||||
if num_images is None:
|
||||
num_images = len(self.image_keys) or len(self.image_features) or MOLMOACT2_DEFAULT_NUM_IMAGES
|
||||
if state_dim is None:
|
||||
state_feature = self.robot_state_feature
|
||||
state_dim = int(state_feature.shape[0]) if state_feature is not None else 0
|
||||
if action_dim is None:
|
||||
action_feature = self.action_feature
|
||||
action_dim = (
|
||||
int(action_feature.shape[0]) if action_feature is not None else self.expected_max_action_dim
|
||||
)
|
||||
if action_horizon is None:
|
||||
action_horizon = self.chunk_size
|
||||
if include_discrete_action is None:
|
||||
include_discrete_action = self.action_mode in {"discrete", "both"}
|
||||
|
||||
return infer_molmoact2_max_sequence_length(
|
||||
num_images=int(num_images),
|
||||
state_dim=int(state_dim),
|
||||
action_dim=int(action_dim),
|
||||
action_horizon=int(action_horizon),
|
||||
include_discrete_action=bool(include_discrete_action),
|
||||
)
|
||||
|
||||
@property
|
||||
def observation_delta_indices(self) -> None:
|
||||
return None
|
||||
|
||||
@property
|
||||
def action_delta_indices(self) -> list[int]:
|
||||
return list(range(self.chunk_size))
|
||||
|
||||
@property
|
||||
def reward_delta_indices(self) -> None:
|
||||
return None
|
||||
|
||||
def get_optimizer_preset(self) -> OptimizerConfig:
|
||||
return AdamWConfig(
|
||||
lr=self.optimizer_lr,
|
||||
betas=self.optimizer_betas,
|
||||
eps=self.optimizer_eps,
|
||||
weight_decay=self.optimizer_weight_decay,
|
||||
grad_clip_norm=self.optimizer_grad_clip_norm,
|
||||
)
|
||||
|
||||
def get_scheduler_preset(self) -> LRSchedulerConfig | None:
|
||||
return MolmoAct2CosineDecayWithWarmupSchedulerConfig(
|
||||
peak_lr=self.optimizer_lr,
|
||||
decay_lr=self.scheduler_decay_lr,
|
||||
num_warmup_steps=self.scheduler_warmup_steps,
|
||||
num_decay_steps=self.scheduler_decay_steps,
|
||||
)
|
||||
|
||||
def set_dataset_feature_metadata(self, features: dict[str, Any]) -> None:
|
||||
self.dataset_feature_names = {}
|
||||
for key in (ACTION, OBS_STATE):
|
||||
feature = features.get(key) if isinstance(features, dict) else None
|
||||
if isinstance(feature, dict) and feature.get("names") is not None:
|
||||
self.dataset_feature_names[key] = feature["names"]
|
||||
|
||||
def validate_features(self) -> None:
|
||||
if OBS_STATE not in self.input_features:
|
||||
self.input_features[OBS_STATE] = PolicyFeature(type=FeatureType.STATE, shape=(0,))
|
||||
if ACTION not in self.output_features:
|
||||
self.output_features[ACTION] = PolicyFeature(type=FeatureType.ACTION, shape=(0,))
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,883 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
from contextlib import suppress
|
||||
from copy import deepcopy
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from huggingface_hub import snapshot_download
|
||||
from torch import Tensor
|
||||
|
||||
from lerobot.configs import PipelineFeatureType, PolicyFeature
|
||||
from lerobot.processor import (
|
||||
AddBatchDimensionProcessorStep,
|
||||
DeviceProcessorStep,
|
||||
NormalizerProcessorStep,
|
||||
PolicyAction,
|
||||
PolicyProcessorPipeline,
|
||||
ProcessorStep,
|
||||
ProcessorStepRegistry,
|
||||
RenameObservationsProcessorStep,
|
||||
UnnormalizerProcessorStep,
|
||||
policy_action_to_transition,
|
||||
transition_to_policy_action,
|
||||
)
|
||||
from lerobot.types import EnvTransition, TransitionKey
|
||||
from lerobot.utils.constants import (
|
||||
ACTION,
|
||||
OBS_IMAGES,
|
||||
OBS_STATE,
|
||||
POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
)
|
||||
from lerobot.utils.import_utils import require_package
|
||||
|
||||
from .configuration_molmoact2 import MolmoAct2Config, infer_molmoact2_max_sequence_length
|
||||
|
||||
ACTION_OUTPUT_TOKEN = "<action_output>" # nosec B105
|
||||
ACTION_START_TOKEN = "<action_start>" # nosec B105
|
||||
ACTION_END_TOKEN = "<action_end>" # nosec B105
|
||||
ACTION_TOKEN_PREFIX = "<action_" # nosec B105
|
||||
STATE_START_TOKEN = "<state_start>" # nosec B105
|
||||
STATE_END_TOKEN = "<state_end>" # nosec B105
|
||||
STATE_TOKEN_PREFIX = "<state_" # nosec B105
|
||||
SETUP_START_TOKEN = "<setup_start>" # nosec B105
|
||||
SETUP_END_TOKEN = "<setup_end>" # nosec B105
|
||||
CONTROL_START_TOKEN = "<control_start>" # nosec B105
|
||||
CONTROL_END_TOKEN = "<control_end>" # nosec B105
|
||||
|
||||
_QUESTION_TRAILING_SENTENCE_PUNCTUATION = ".,!?;:,\u2026"
|
||||
_QUESTION_TRAILING_CLOSERS = "\"'\u201d\u2019)]}"
|
||||
_QUESTION_SURROUNDING_DELIMITERS = "\"'`\u201c\u201d\u2018\u2019[](){}"
|
||||
_QUESTION_PREFIX_PATTERNS = tuple(
|
||||
re.compile(pattern, flags=re.IGNORECASE)
|
||||
for pattern in (
|
||||
r"^(?:task|instruction|language[_ ]instruction|goal)\s*[:\-]\s*",
|
||||
r"^(?:the\s+task\s+is\s+to|your\s+task\s+is\s+to)\s+",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _hf_token() -> str | None:
|
||||
return os.environ.get("HF_TOKEN") or os.environ.get("HF_ACCESS_TOKEN")
|
||||
|
||||
|
||||
def _resolve_checkpoint_location(
|
||||
checkpoint_path: str,
|
||||
*,
|
||||
revision: str | None = None,
|
||||
force_download: bool = False,
|
||||
) -> str:
|
||||
checkpoint_path = str(checkpoint_path or "").strip()
|
||||
if not checkpoint_path:
|
||||
raise ValueError("MolmoAct2 policy requires `checkpoint_path`.")
|
||||
local_path = Path(checkpoint_path).expanduser()
|
||||
if local_path.exists():
|
||||
return str(local_path)
|
||||
return snapshot_download(
|
||||
repo_id=checkpoint_path,
|
||||
repo_type="model",
|
||||
revision=revision,
|
||||
force_download=force_download,
|
||||
token=_hf_token(),
|
||||
)
|
||||
|
||||
|
||||
def _load_hf_norm_stats_for_tag(
|
||||
checkpoint_path: str,
|
||||
*,
|
||||
revision: str | None,
|
||||
force_download: bool,
|
||||
norm_tag: str | None,
|
||||
) -> tuple[dict[str, dict[str, Any]], dict[str, Any]]:
|
||||
norm_tag = str(norm_tag or "").strip()
|
||||
if not norm_tag:
|
||||
raise ValueError("MolmoAct2 HF checkpoint inference requires `policy.norm_tag` for normalization.")
|
||||
|
||||
checkpoint_location = Path(
|
||||
_resolve_checkpoint_location(
|
||||
checkpoint_path,
|
||||
revision=revision,
|
||||
force_download=force_download,
|
||||
)
|
||||
)
|
||||
config_path = checkpoint_location / "config.json"
|
||||
norm_stats_filename = "norm_stats.json"
|
||||
if config_path.exists():
|
||||
with suppress(OSError, json.JSONDecodeError):
|
||||
norm_stats_filename = str(
|
||||
json.loads(config_path.read_text()).get("norm_stats_filename") or norm_stats_filename
|
||||
)
|
||||
|
||||
stats_path = checkpoint_location / norm_stats_filename
|
||||
if not stats_path.exists():
|
||||
raise FileNotFoundError(
|
||||
f"MolmoAct2 HF checkpoint is missing {norm_stats_filename!r}; cannot resolve norm_tag={norm_tag!r}."
|
||||
)
|
||||
payload = json.loads(stats_path.read_text())
|
||||
metadata_by_tag = payload.get("metadata_by_tag")
|
||||
if not isinstance(metadata_by_tag, dict):
|
||||
raise ValueError(f"MolmoAct2 norm stats file {stats_path} has no metadata_by_tag mapping.")
|
||||
metadata = metadata_by_tag.get(norm_tag)
|
||||
if metadata is None:
|
||||
available = sorted(str(tag) for tag in metadata_by_tag)
|
||||
raise ValueError(f"Unknown MolmoAct2 norm_tag={norm_tag!r}. Available tags: {available}.")
|
||||
if not isinstance(metadata, dict):
|
||||
raise ValueError(f"MolmoAct2 norm_tag={norm_tag!r} metadata must be a mapping.")
|
||||
|
||||
def numeric_stats(raw_stats: dict[str, Any]) -> dict[str, Any]:
|
||||
stats: dict[str, Any] = {}
|
||||
for key, value in raw_stats.items():
|
||||
if key == "names":
|
||||
continue
|
||||
if isinstance(value, (list, tuple)) and any(isinstance(item, str) for item in value):
|
||||
continue
|
||||
stats[key] = deepcopy(value)
|
||||
return stats
|
||||
|
||||
action_stats = metadata.get("action_stats")
|
||||
state_stats = metadata.get("state_stats")
|
||||
if not isinstance(action_stats, dict) or not isinstance(state_stats, dict):
|
||||
raise ValueError(f"MolmoAct2 norm_tag={norm_tag!r} must define action_stats and state_stats.")
|
||||
return {ACTION: numeric_stats(action_stats), OBS_STATE: numeric_stats(state_stats)}, metadata
|
||||
|
||||
|
||||
def _to_numpy(value: Any) -> np.ndarray:
|
||||
if isinstance(value, np.ndarray):
|
||||
return value
|
||||
if torch.is_tensor(value):
|
||||
return value.detach().cpu().numpy()
|
||||
return np.asarray(value)
|
||||
|
||||
|
||||
def _normalize_image(value: Any) -> np.ndarray:
|
||||
arr = _to_numpy(value)
|
||||
while arr.ndim > 3 and int(arr.shape[0]) == 1:
|
||||
arr = arr[0]
|
||||
if arr.ndim == 2:
|
||||
arr = np.stack([arr] * 3, axis=-1)
|
||||
if arr.ndim == 3 and arr.shape[0] in {1, 3, 4} and arr.shape[-1] not in {1, 3, 4}:
|
||||
arr = np.moveaxis(arr, 0, -1)
|
||||
if arr.ndim == 3 and arr.shape[-1] == 1:
|
||||
arr = np.repeat(arr, 3, axis=-1)
|
||||
if arr.ndim != 3 or arr.shape[-1] not in {3, 4}:
|
||||
raise ValueError(f"Unsupported image shape for MolmoAct2: {arr.shape}.")
|
||||
if arr.shape[-1] == 4:
|
||||
arr = arr[..., :3]
|
||||
if arr.dtype in (np.float16, np.float32, np.float64):
|
||||
if arr.size > 0 and float(np.nanmax(arr)) <= 1.0:
|
||||
arr = arr * 255.0
|
||||
arr = np.clip(arr, 0, 255).astype(np.uint8)
|
||||
elif arr.dtype != np.uint8:
|
||||
arr = np.clip(arr, 0, 255).astype(np.uint8)
|
||||
return arr
|
||||
|
||||
|
||||
def _normalize_question_text(text: str) -> str:
|
||||
normalized = re.sub(r"\s+", " ", str(text or "")).strip()
|
||||
if not normalized:
|
||||
return ""
|
||||
previous = None
|
||||
while normalized and normalized != previous:
|
||||
previous = normalized
|
||||
normalized = normalized.strip().strip(_QUESTION_SURROUNDING_DELIMITERS).strip()
|
||||
for pattern in _QUESTION_PREFIX_PATTERNS:
|
||||
normalized = pattern.sub("", normalized, count=1).strip()
|
||||
normalized = normalized.rstrip(_QUESTION_TRAILING_SENTENCE_PUNCTUATION).rstrip()
|
||||
normalized = normalized.rstrip(_QUESTION_TRAILING_CLOSERS).rstrip()
|
||||
normalized = normalized.rstrip(_QUESTION_TRAILING_SENTENCE_PUNCTUATION).rstrip()
|
||||
chunks = [chunk.strip() for chunk in re.split(r"[.!?]+", normalized) if chunk.strip()]
|
||||
if len(chunks) > 1:
|
||||
normalized = "; ".join(chunks)
|
||||
return normalized.lower()
|
||||
|
||||
|
||||
def _wrap_setup_text(setup_type: str, add_setup_tokens: bool) -> str:
|
||||
setup_type = str(setup_type or "")
|
||||
if setup_type.startswith(SETUP_START_TOKEN) and setup_type.endswith(SETUP_END_TOKEN):
|
||||
return setup_type
|
||||
if not setup_type or not add_setup_tokens:
|
||||
return setup_type
|
||||
return f"{SETUP_START_TOKEN}{setup_type}{SETUP_END_TOKEN}"
|
||||
|
||||
|
||||
def _wrap_control_text(control_mode: str, add_control_tokens: bool) -> str:
|
||||
control_mode = str(control_mode or "")
|
||||
if control_mode.startswith(CONTROL_START_TOKEN) and control_mode.endswith(CONTROL_END_TOKEN):
|
||||
return control_mode
|
||||
if not control_mode or not add_control_tokens:
|
||||
return control_mode
|
||||
return f"{CONTROL_START_TOKEN}{control_mode}{CONTROL_END_TOKEN}"
|
||||
|
||||
|
||||
def _build_discrete_state_string(state: np.ndarray, num_state_tokens: int) -> str:
|
||||
if num_state_tokens <= 0:
|
||||
raise ValueError(f"num_state_tokens must be > 0, got {num_state_tokens}.")
|
||||
arr = np.asarray(state, dtype=np.float32)
|
||||
arr = np.nan_to_num(arr, nan=0.0, posinf=1.0, neginf=-1.0)
|
||||
arr = np.clip(arr, -1.0, 1.0)
|
||||
scaled = (arr + 1.0) / 2.0 * float(num_state_tokens - 1)
|
||||
token_ids = np.clip(np.rint(scaled).astype(np.int64), 0, int(num_state_tokens) - 1).reshape(-1)
|
||||
return f"{STATE_START_TOKEN}{''.join(f'{STATE_TOKEN_PREFIX}{int(token_id)}>' for token_id in token_ids)}{STATE_END_TOKEN}"
|
||||
|
||||
|
||||
def _build_robot_text(
|
||||
*,
|
||||
task: str,
|
||||
discrete_state_string: str,
|
||||
setup_type: str,
|
||||
control_mode: str,
|
||||
add_setup_tokens: bool,
|
||||
add_control_tokens: bool,
|
||||
num_images: int,
|
||||
) -> str:
|
||||
setup_text = _wrap_setup_text(setup_type, add_setup_tokens=add_setup_tokens)
|
||||
control_text = _wrap_control_text(control_mode, add_control_tokens=add_control_tokens)
|
||||
state_clause = (
|
||||
f" The current state of the robot is {discrete_state_string}." if discrete_state_string else ""
|
||||
)
|
||||
prompt = (
|
||||
f"The task is to {task}. The setup is {setup_text}.{state_clause} "
|
||||
f"The expected control mode is {control_text}. Given these, what action should the robot take to complete the task?"
|
||||
)
|
||||
if num_images <= 0:
|
||||
image_prefix = ""
|
||||
elif num_images == 1:
|
||||
image_prefix = "<|image|>"
|
||||
else:
|
||||
image_prefix = "".join(f"Image {idx + 1}<|image|>" for idx in range(num_images))
|
||||
return f"{image_prefix}<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n{ACTION_OUTPUT_TOKEN}"
|
||||
|
||||
|
||||
def _as_text_list(value: Any, batch_size: int) -> list[str]:
|
||||
if value is None:
|
||||
return [""] * batch_size
|
||||
if isinstance(value, str):
|
||||
return [value] * batch_size
|
||||
if torch.is_tensor(value):
|
||||
if value.ndim == 0:
|
||||
return [str(value.item())] * batch_size
|
||||
flat = value.detach().cpu().reshape(-1).tolist()
|
||||
texts = [str(item) for item in flat]
|
||||
elif isinstance(value, np.ndarray):
|
||||
if value.ndim == 0:
|
||||
return [str(value.item())] * batch_size
|
||||
texts = [str(item) for item in value.reshape(-1).tolist()]
|
||||
elif isinstance(value, (list, tuple)):
|
||||
texts = [str(item) for item in value]
|
||||
else:
|
||||
texts = [str(value)]
|
||||
if len(texts) == batch_size:
|
||||
return texts
|
||||
if len(texts) == 1:
|
||||
return texts * batch_size
|
||||
raise ValueError(f"Expected {batch_size} task strings, got {len(texts)}.")
|
||||
|
||||
|
||||
def _tokenize_discrete_action(action: np.ndarray, processor: Any) -> list[int]:
|
||||
arr = np.asarray(action, dtype=np.float32)
|
||||
if arr.ndim == 2:
|
||||
arr = arr[None, :, :]
|
||||
elif arr.ndim == 1:
|
||||
arr = arr[None, None, :]
|
||||
tokens_out = processor(arr)
|
||||
if isinstance(tokens_out, dict):
|
||||
tokens_out = tokens_out.get("input_ids", next(iter(tokens_out.values())))
|
||||
if isinstance(tokens_out, np.ndarray):
|
||||
tokens_out = tokens_out.tolist()
|
||||
if torch.is_tensor(tokens_out):
|
||||
tokens_out = tokens_out.detach().cpu().tolist()
|
||||
if not isinstance(tokens_out, list):
|
||||
raise TypeError(f"Unexpected discrete action tokenizer output type: {type(tokens_out)}")
|
||||
if tokens_out and isinstance(tokens_out[0], (list, tuple, np.ndarray)):
|
||||
tokens_out = tokens_out[0]
|
||||
return [int(token_id) for token_id in tokens_out]
|
||||
|
||||
|
||||
def _build_discrete_action_string(action: np.ndarray, processor: Any) -> str:
|
||||
token_ids = _tokenize_discrete_action(action, processor)
|
||||
pieces = "".join(f"{ACTION_TOKEN_PREFIX}{int(token_id)}>" for token_id in token_ids)
|
||||
return f"{ACTION_START_TOKEN}{pieces}{ACTION_END_TOKEN}"
|
||||
|
||||
|
||||
def _single_token_id(tokenizer: Any, token: str) -> int:
|
||||
token_ids = tokenizer.encode(token, add_special_tokens=False)
|
||||
if len(token_ids) != 1:
|
||||
raise ValueError(f"MolmoAct2 token {token!r} must encode to one token, got {token_ids}.")
|
||||
return int(token_ids[0])
|
||||
|
||||
|
||||
def _flatten_feature_names(raw_names: Any) -> list[str] | None:
|
||||
if raw_names is None:
|
||||
return None
|
||||
if isinstance(raw_names, dict):
|
||||
names: list[str] = []
|
||||
for value in raw_names.values():
|
||||
if isinstance(value, (list, tuple)):
|
||||
names.extend(str(item) for item in value)
|
||||
elif value is not None:
|
||||
names.append(str(value))
|
||||
return names or None
|
||||
if isinstance(raw_names, (list, tuple)):
|
||||
names = [str(item) for item in raw_names]
|
||||
return names or None
|
||||
return [str(raw_names)]
|
||||
|
||||
|
||||
def _feature_dim(stats: dict[str, Any] | None) -> int | None:
|
||||
if not isinstance(stats, dict):
|
||||
return None
|
||||
for key in ("mean", "std", "min", "max", "q01", "q99", "q10", "q90", "mask"):
|
||||
value = stats.get(key)
|
||||
if value is None:
|
||||
continue
|
||||
if torch.is_tensor(value):
|
||||
return int(value.shape[-1]) if value.ndim > 0 else None
|
||||
arr = np.asarray(value)
|
||||
return int(arr.shape[-1]) if arr.ndim > 0 else None
|
||||
return None
|
||||
|
||||
|
||||
def _feature_names_from_meta(dataset_meta: Any | None, feature_key: str) -> list[str] | None:
|
||||
if dataset_meta is None:
|
||||
return None
|
||||
|
||||
root = getattr(dataset_meta, "root", None)
|
||||
candidate_roots = []
|
||||
if root is not None:
|
||||
repo_id = str(getattr(dataset_meta, "repo_id", "") or "").strip()
|
||||
if repo_id:
|
||||
candidate_roots.append(Path(root) / repo_id)
|
||||
candidate_roots.append(Path(root))
|
||||
for candidate_root in candidate_roots:
|
||||
info_path = candidate_root / "meta" / "info.json"
|
||||
if info_path.exists():
|
||||
try:
|
||||
with info_path.open("r", encoding="utf-8") as f:
|
||||
info = json.load(f)
|
||||
names = _flatten_feature_names((info.get("features") or {}).get(feature_key, {}).get("names"))
|
||||
if names:
|
||||
return names
|
||||
except (OSError, json.JSONDecodeError, AttributeError):
|
||||
pass
|
||||
|
||||
for container in (
|
||||
getattr(getattr(dataset_meta, "info", None), "features", None),
|
||||
getattr(dataset_meta, "features", None),
|
||||
):
|
||||
if not isinstance(container, dict):
|
||||
continue
|
||||
feature = container.get(feature_key)
|
||||
if not isinstance(feature, dict):
|
||||
continue
|
||||
names = _flatten_feature_names(feature.get("names"))
|
||||
if names:
|
||||
return names
|
||||
return None
|
||||
|
||||
|
||||
def _add_gripper_masks_to_stats(
|
||||
dataset_stats: dict[str, dict[str, Any]] | None,
|
||||
dataset_meta: Any | None,
|
||||
*,
|
||||
normalize_gripper: bool,
|
||||
dataset_feature_names: dict[str, Any] | None = None,
|
||||
) -> dict[str, dict[str, Any]] | None:
|
||||
if not dataset_stats:
|
||||
return dataset_stats
|
||||
|
||||
stats = deepcopy(dataset_stats)
|
||||
for key in (ACTION, OBS_STATE):
|
||||
feature_stats = stats.get(key)
|
||||
if not isinstance(feature_stats, dict):
|
||||
continue
|
||||
dim = _feature_dim(feature_stats)
|
||||
if dim is None:
|
||||
continue
|
||||
|
||||
if normalize_gripper:
|
||||
feature_stats["mask"] = [True] * dim
|
||||
continue
|
||||
|
||||
names = _flatten_feature_names((dataset_feature_names or {}).get(key))
|
||||
if names is None:
|
||||
names = _feature_names_from_meta(dataset_meta, key)
|
||||
if names is None:
|
||||
names = _flatten_feature_names(feature_stats.get("names"))
|
||||
if names is None:
|
||||
continue
|
||||
if len(names) != dim:
|
||||
continue
|
||||
feature_stats["mask"] = ["gripper" not in name.lower() for name in names]
|
||||
return stats
|
||||
|
||||
|
||||
class _MolmoAct2MaskedNormalizationMixin:
|
||||
def _apply_transform(
|
||||
self, tensor: Tensor, key: str, feature_type: Any, *, inverse: bool = False
|
||||
) -> Tensor:
|
||||
transformed = super()._apply_transform(tensor, key, feature_type, inverse=inverse)
|
||||
stats = getattr(self, "_tensor_stats", {}).get(key, {})
|
||||
mask = stats.get("mask") if isinstance(stats, dict) else None
|
||||
if mask is None:
|
||||
return transformed
|
||||
mask = mask.to(device=tensor.device, dtype=torch.bool)
|
||||
if mask.ndim != 1 or tensor.shape[-1] != mask.shape[0]:
|
||||
return transformed
|
||||
while mask.ndim < tensor.ndim:
|
||||
mask = mask.unsqueeze(0)
|
||||
return torch.where(mask, transformed, tensor)
|
||||
|
||||
|
||||
@ProcessorStepRegistry.register(name="molmoact2_masked_normalizer")
|
||||
@dataclass
|
||||
class MolmoAct2MaskedNormalizerProcessorStep(_MolmoAct2MaskedNormalizationMixin, NormalizerProcessorStep):
|
||||
pass
|
||||
|
||||
|
||||
@ProcessorStepRegistry.register(name="molmoact2_masked_unnormalizer")
|
||||
@dataclass
|
||||
class MolmoAct2MaskedUnnormalizerProcessorStep(_MolmoAct2MaskedNormalizationMixin, UnnormalizerProcessorStep):
|
||||
pass
|
||||
|
||||
|
||||
@ProcessorStepRegistry.register(name="molmoact2_clamp_normalized")
|
||||
@dataclass
|
||||
class MolmoAct2ClampNormalizedProcessorStep(ProcessorStep):
|
||||
"""Clamp q01/q99-normalized state and action to the range used by the old trainer."""
|
||||
|
||||
def __call__(self, transition: EnvTransition) -> EnvTransition:
|
||||
transition = transition.copy()
|
||||
observation = transition.get(TransitionKey.OBSERVATION)
|
||||
if isinstance(observation, dict) and OBS_STATE in observation:
|
||||
observation = observation.copy()
|
||||
observation[OBS_STATE] = torch.as_tensor(observation[OBS_STATE]).clamp(-1.0, 1.0)
|
||||
transition[TransitionKey.OBSERVATION] = observation
|
||||
action = transition.get(TransitionKey.ACTION)
|
||||
if action is not None:
|
||||
transition[TransitionKey.ACTION] = torch.as_tensor(action).clamp(-1.0, 1.0)
|
||||
return transition
|
||||
|
||||
def transform_features(
|
||||
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
|
||||
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
|
||||
return features
|
||||
|
||||
|
||||
@ProcessorStepRegistry.register(name="molmoact2_pack_inputs")
|
||||
@dataclass
|
||||
class MolmoAct2PackInputsProcessorStep(ProcessorStep):
|
||||
checkpoint_path: str
|
||||
checkpoint_revision: str | None = None
|
||||
checkpoint_force_download: bool = False
|
||||
trust_remote_code: bool = True
|
||||
action_mode: str = "both"
|
||||
discrete_action_tokenizer: str = "allenai/MolmoAct2-FAST-Tokenizer"
|
||||
image_keys: list[str] = field(default_factory=list)
|
||||
setup_type: str = ""
|
||||
control_mode: str = ""
|
||||
normalize_language: bool = True
|
||||
add_setup_tokens: bool = True
|
||||
add_control_tokens: bool = True
|
||||
num_state_tokens: int = 256
|
||||
max_sequence_length: int | None = None
|
||||
chunk_size: int = 30
|
||||
max_action_dim: int = 32
|
||||
env_action_dim: int | None = None
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
require_package("transformers", extra="molmoact2")
|
||||
from transformers import AutoProcessor
|
||||
|
||||
checkpoint_location = _resolve_checkpoint_location(
|
||||
self.checkpoint_path,
|
||||
revision=self.checkpoint_revision,
|
||||
force_download=bool(self.checkpoint_force_download),
|
||||
)
|
||||
self.processor = AutoProcessor.from_pretrained(
|
||||
checkpoint_location,
|
||||
trust_remote_code=self.trust_remote_code,
|
||||
use_fast=False,
|
||||
token=_hf_token(),
|
||||
)
|
||||
self.action_processor = None
|
||||
if self.action_mode in {"discrete", "both"}:
|
||||
self.action_processor = AutoProcessor.from_pretrained(
|
||||
self.discrete_action_tokenizer,
|
||||
trust_remote_code=self.trust_remote_code,
|
||||
token=_hf_token(),
|
||||
)
|
||||
self._action_start_id = _single_token_id(self.processor.tokenizer, ACTION_START_TOKEN)
|
||||
self._action_end_id = _single_token_id(self.processor.tokenizer, ACTION_END_TOKEN)
|
||||
self._eos_token = self.processor.tokenizer.eos_token or ""
|
||||
self._eos_token_id = self.processor.tokenizer.eos_token_id
|
||||
|
||||
def get_config(self) -> dict[str, Any]:
|
||||
return {
|
||||
"checkpoint_path": self.checkpoint_path,
|
||||
"checkpoint_revision": self.checkpoint_revision,
|
||||
"checkpoint_force_download": self.checkpoint_force_download,
|
||||
"trust_remote_code": self.trust_remote_code,
|
||||
"action_mode": self.action_mode,
|
||||
"discrete_action_tokenizer": self.discrete_action_tokenizer,
|
||||
"image_keys": list(self.image_keys),
|
||||
"setup_type": self.setup_type,
|
||||
"control_mode": self.control_mode,
|
||||
"normalize_language": self.normalize_language,
|
||||
"add_setup_tokens": self.add_setup_tokens,
|
||||
"add_control_tokens": self.add_control_tokens,
|
||||
"num_state_tokens": self.num_state_tokens,
|
||||
"max_sequence_length": self.max_sequence_length,
|
||||
"chunk_size": self.chunk_size,
|
||||
"max_action_dim": self.max_action_dim,
|
||||
"env_action_dim": self.env_action_dim,
|
||||
}
|
||||
|
||||
def _resolve_max_sequence_length(
|
||||
self,
|
||||
*,
|
||||
num_images: int,
|
||||
state_dim: int,
|
||||
action_dim: int,
|
||||
action_horizon: int,
|
||||
include_discrete_action: bool,
|
||||
) -> int:
|
||||
if self.max_sequence_length is not None:
|
||||
return int(self.max_sequence_length)
|
||||
return infer_molmoact2_max_sequence_length(
|
||||
num_images=num_images,
|
||||
state_dim=state_dim,
|
||||
action_dim=action_dim,
|
||||
action_horizon=action_horizon,
|
||||
include_discrete_action=include_discrete_action,
|
||||
)
|
||||
|
||||
def _batch_size(self, observation: dict[str, Any], action: Tensor | None) -> int:
|
||||
if action is not None:
|
||||
return int(action.shape[0])
|
||||
state = observation.get(OBS_STATE)
|
||||
if torch.is_tensor(state) or isinstance(state, np.ndarray):
|
||||
return int(state.shape[0]) if getattr(state, "ndim", 0) > 1 else 1
|
||||
for key in self._resolve_image_keys(observation):
|
||||
value = observation[key]
|
||||
if torch.is_tensor(value) or isinstance(value, np.ndarray):
|
||||
return int(value.shape[0]) if getattr(value, "ndim", 0) == 4 else 1
|
||||
return 1
|
||||
|
||||
def _resolve_image_keys(self, observation: dict[str, Any]) -> list[str]:
|
||||
if self.image_keys:
|
||||
missing = [key for key in self.image_keys if key not in observation]
|
||||
if missing:
|
||||
raise ValueError(f"MolmoAct2 image_keys missing from observation: {missing}.")
|
||||
return list(self.image_keys)
|
||||
keys = [key for key in observation if str(key).startswith(f"{OBS_IMAGES}.")]
|
||||
if not keys:
|
||||
keys = [key for key in observation if str(key).startswith("observation.image")]
|
||||
if not keys:
|
||||
raise ValueError("MolmoAct2 requires at least one image observation.")
|
||||
return sorted(keys)
|
||||
|
||||
def _extract_images(self, observation: dict[str, Any], batch_size: int) -> list[list[np.ndarray]]:
|
||||
images_by_example: list[list[np.ndarray]] = [[] for _ in range(batch_size)]
|
||||
for key in self._resolve_image_keys(observation):
|
||||
value = observation[key]
|
||||
for batch_idx in range(batch_size):
|
||||
item = value
|
||||
if (torch.is_tensor(value) or isinstance(value, np.ndarray)) and getattr(
|
||||
value, "ndim", 0
|
||||
) >= 4:
|
||||
item = value[batch_idx]
|
||||
images_by_example[batch_idx].append(_normalize_image(item))
|
||||
return images_by_example
|
||||
|
||||
def _extract_state(self, observation: dict[str, Any], batch_size: int) -> Tensor:
|
||||
if OBS_STATE not in observation:
|
||||
raise ValueError("MolmoAct2 requires observation.state for discrete state prompting.")
|
||||
state = torch.as_tensor(observation[OBS_STATE], dtype=torch.float32)
|
||||
if state.ndim == 1:
|
||||
state = state.unsqueeze(0)
|
||||
if int(state.shape[0]) != batch_size:
|
||||
raise ValueError(f"State batch size {state.shape[0]} does not match batch size {batch_size}.")
|
||||
return state
|
||||
|
||||
def _pad_action(self, action: Tensor, action_is_pad: Any | None) -> tuple[Tensor, Tensor, Tensor]:
|
||||
if action.ndim == 2:
|
||||
action = action.unsqueeze(1)
|
||||
if action.ndim != 3:
|
||||
raise ValueError(f"MolmoAct2 expected action shape [B, T, D], got {tuple(action.shape)}.")
|
||||
if action.shape[-1] > self.max_action_dim:
|
||||
raise ValueError(
|
||||
f"Action dim {action.shape[-1]} exceeds MolmoAct2 max_action_dim={self.max_action_dim}."
|
||||
)
|
||||
padded = torch.zeros(
|
||||
(*action.shape[:-1], self.max_action_dim),
|
||||
device=action.device,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
padded[..., : action.shape[-1]] = action.to(dtype=torch.float32)
|
||||
action_dim_is_pad = torch.ones(
|
||||
(action.shape[0], self.max_action_dim), device=action.device, dtype=torch.bool
|
||||
)
|
||||
action_dim_is_pad[:, : action.shape[-1]] = False
|
||||
if action_is_pad is None:
|
||||
action_horizon_is_pad = torch.zeros(action.shape[:2], device=action.device, dtype=torch.bool)
|
||||
else:
|
||||
action_horizon_is_pad = torch.as_tensor(action_is_pad, device=action.device, dtype=torch.bool)
|
||||
if action_horizon_is_pad.ndim == 1:
|
||||
action_horizon_is_pad = action_horizon_is_pad.unsqueeze(0)
|
||||
if tuple(action_horizon_is_pad.shape) != tuple(action.shape[:2]):
|
||||
raise ValueError(
|
||||
"action_is_pad must match action horizon shape: "
|
||||
f"got {tuple(action_horizon_is_pad.shape)} for action {tuple(action.shape)}."
|
||||
)
|
||||
return padded, action_horizon_is_pad, action_dim_is_pad
|
||||
|
||||
def _build_labels(self, input_ids: Tensor, attention_mask: Tensor) -> Tensor:
|
||||
labels = torch.full_like(input_ids, -100)
|
||||
for batch_idx in range(input_ids.shape[0]):
|
||||
valid = attention_mask[batch_idx].to(dtype=torch.bool)
|
||||
row = input_ids[batch_idx]
|
||||
starts = (row == self._action_start_id).nonzero(as_tuple=False).flatten().tolist()
|
||||
ends = (row == self._action_end_id).nonzero(as_tuple=False).flatten().tolist()
|
||||
end_ptr = 0
|
||||
for start in starts:
|
||||
while end_ptr < len(ends) and ends[end_ptr] < start:
|
||||
end_ptr += 1
|
||||
if end_ptr >= len(ends):
|
||||
raise ValueError(
|
||||
"Found <action_start> without matching <action_end> in MolmoAct2 labels."
|
||||
)
|
||||
end = int(ends[end_ptr])
|
||||
label_end = end + 1
|
||||
if (
|
||||
self._eos_token_id is not None
|
||||
and label_end < int(row.shape[0])
|
||||
and int(row[label_end]) == int(self._eos_token_id)
|
||||
):
|
||||
label_end += 1
|
||||
labels[batch_idx, start:label_end] = row[start:label_end]
|
||||
end_ptr += 1
|
||||
if not starts:
|
||||
raise ValueError("No discrete action span found in MolmoAct2 training text.")
|
||||
labels[batch_idx] = torch.where(
|
||||
valid, labels[batch_idx], torch.full_like(labels[batch_idx], -100)
|
||||
)
|
||||
return labels
|
||||
|
||||
def __call__(self, transition: EnvTransition) -> EnvTransition:
|
||||
transition = transition.copy()
|
||||
observation = transition.get(TransitionKey.OBSERVATION) or {}
|
||||
if not isinstance(observation, dict):
|
||||
raise ValueError("MolmoAct2 expected an observation dictionary.")
|
||||
complementary = dict(transition.get(TransitionKey.COMPLEMENTARY_DATA) or {})
|
||||
|
||||
raw_action = transition.get(TransitionKey.ACTION)
|
||||
action = torch.as_tensor(raw_action, dtype=torch.float32) if raw_action is not None else None
|
||||
batch_size = self._batch_size(observation, action)
|
||||
state = self._extract_state(observation, batch_size)
|
||||
images_by_example = self._extract_images(observation, batch_size)
|
||||
|
||||
task_source = complementary.get("task")
|
||||
if task_source is None:
|
||||
task_source = observation.get("task")
|
||||
if task_source is None:
|
||||
task_source = observation.get("observation.language")
|
||||
if task_source is None:
|
||||
task_source = complementary.get("language_instruction")
|
||||
tasks = _as_text_list(task_source, batch_size)
|
||||
if self.normalize_language:
|
||||
tasks = [_normalize_question_text(task) for task in tasks]
|
||||
complementary["task"] = tasks
|
||||
|
||||
action_padded = None
|
||||
action_horizon_is_pad = None
|
||||
action_dim_is_pad = torch.ones((batch_size, self.max_action_dim), dtype=torch.bool)
|
||||
real_action_dim = int(self.env_action_dim or 0)
|
||||
if action is not None:
|
||||
action_is_pad = complementary.get("action_is_pad")
|
||||
if action_is_pad is None:
|
||||
action_is_pad = complementary.get("action_horizon_is_pad")
|
||||
action_padded, action_horizon_is_pad, action_dim_is_pad = self._pad_action(action, action_is_pad)
|
||||
real_action_dim = int(action.shape[-1])
|
||||
elif real_action_dim > 0:
|
||||
action_dim_is_pad[:, :real_action_dim] = False
|
||||
|
||||
prompt_texts: list[str] = []
|
||||
full_texts: list[str] = []
|
||||
flat_images: list[np.ndarray] = []
|
||||
state_np = state.detach().cpu().numpy()
|
||||
build_action_labels = action is not None and self.action_mode in {"discrete", "both"}
|
||||
for batch_idx in range(batch_size):
|
||||
images = images_by_example[batch_idx]
|
||||
flat_images.extend(images)
|
||||
discrete_state = _build_discrete_state_string(state_np[batch_idx], self.num_state_tokens)
|
||||
prompt = _build_robot_text(
|
||||
task=tasks[batch_idx],
|
||||
discrete_state_string=discrete_state,
|
||||
setup_type=self.setup_type,
|
||||
control_mode=self.control_mode,
|
||||
add_setup_tokens=self.add_setup_tokens,
|
||||
add_control_tokens=self.add_control_tokens,
|
||||
num_images=len(images),
|
||||
)
|
||||
prompt_texts.append(prompt)
|
||||
if build_action_labels:
|
||||
if self.action_processor is None:
|
||||
raise ValueError("Discrete MolmoAct2 training requires an action tokenizer.")
|
||||
answer = _build_discrete_action_string(
|
||||
action[batch_idx].detach().cpu().numpy(), self.action_processor
|
||||
)
|
||||
full_texts.append(f"{prompt}{answer}{self._eos_token}")
|
||||
else:
|
||||
full_texts.append(prompt)
|
||||
|
||||
text = full_texts if build_action_labels else prompt_texts
|
||||
inputs = self.processor(text=text, images=flat_images, return_tensors="pt", padding=True)
|
||||
if action is None:
|
||||
action_horizon = self.chunk_size
|
||||
elif action.ndim == 2:
|
||||
action_horizon = 1
|
||||
else:
|
||||
action_horizon = int(action.shape[1])
|
||||
max_sequence_length = self._resolve_max_sequence_length(
|
||||
num_images=max((len(images) for images in images_by_example), default=0),
|
||||
state_dim=int(state.shape[-1]),
|
||||
action_dim=max(real_action_dim, 1),
|
||||
action_horizon=action_horizon,
|
||||
include_discrete_action=build_action_labels,
|
||||
)
|
||||
if int(inputs["input_ids"].shape[1]) > max_sequence_length:
|
||||
raise ValueError(
|
||||
f"MolmoAct2 sequence length {int(inputs['input_ids'].shape[1])} exceeds "
|
||||
f"max_sequence_length={max_sequence_length}."
|
||||
)
|
||||
|
||||
if build_action_labels:
|
||||
inputs["labels"] = self._build_labels(inputs["input_ids"], inputs["attention_mask"])
|
||||
|
||||
complementary.update(dict(inputs))
|
||||
complementary["action_dim_is_pad"] = action_dim_is_pad
|
||||
if action_horizon_is_pad is not None:
|
||||
complementary["action_horizon_is_pad"] = action_horizon_is_pad
|
||||
|
||||
if action_padded is not None:
|
||||
transition[TransitionKey.ACTION] = action_padded
|
||||
transition[TransitionKey.COMPLEMENTARY_DATA] = complementary
|
||||
return transition
|
||||
|
||||
def transform_features(
|
||||
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
|
||||
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
|
||||
return features
|
||||
|
||||
|
||||
@ProcessorStepRegistry.register(name="molmoact2_clamp_action")
|
||||
@dataclass
|
||||
class MolmoAct2ClampActionProcessorStep(ProcessorStep):
|
||||
def __call__(self, transition: EnvTransition) -> EnvTransition:
|
||||
transition = transition.copy()
|
||||
action = transition.get(TransitionKey.ACTION)
|
||||
if action is not None:
|
||||
transition[TransitionKey.ACTION] = torch.as_tensor(action).clamp(-1.0, 1.0)
|
||||
return transition
|
||||
|
||||
def transform_features(
|
||||
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
|
||||
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
|
||||
return features
|
||||
|
||||
|
||||
def make_molmoact2_pre_post_processors(
|
||||
config: MolmoAct2Config,
|
||||
dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
|
||||
dataset_meta: Any | None = None,
|
||||
) -> tuple[
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction],
|
||||
]:
|
||||
env_action_dim = None
|
||||
if config.output_features and ACTION in config.output_features:
|
||||
env_action_dim = int(config.output_features[ACTION].shape[0])
|
||||
|
||||
hf_metadata: dict[str, Any] = {}
|
||||
if dataset_stats is None and str(config.norm_tag or "").strip():
|
||||
dataset_stats, hf_metadata = _load_hf_norm_stats_for_tag(
|
||||
config.checkpoint_path,
|
||||
revision=config.checkpoint_revision,
|
||||
force_download=bool(config.checkpoint_force_download),
|
||||
norm_tag=config.norm_tag,
|
||||
)
|
||||
|
||||
image_keys = list(config.image_keys)
|
||||
if not image_keys and isinstance(hf_metadata.get("camera_keys"), list):
|
||||
image_keys = [str(key) for key in hf_metadata["camera_keys"]]
|
||||
setup_type = config.setup_type or str(hf_metadata.get("setup_type") or "")
|
||||
control_mode = config.control_mode or str(hf_metadata.get("control_mode") or "")
|
||||
chunk_size = int(hf_metadata.get("action_horizon") or config.chunk_size)
|
||||
|
||||
masked_dataset_stats = _add_gripper_masks_to_stats(
|
||||
dataset_stats,
|
||||
dataset_meta,
|
||||
normalize_gripper=config.normalize_gripper,
|
||||
dataset_feature_names=config.dataset_feature_names,
|
||||
)
|
||||
|
||||
input_steps: list[ProcessorStep] = [
|
||||
RenameObservationsProcessorStep(rename_map={}),
|
||||
AddBatchDimensionProcessorStep(),
|
||||
MolmoAct2MaskedNormalizerProcessorStep(
|
||||
features={**config.input_features, **config.output_features},
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=masked_dataset_stats,
|
||||
),
|
||||
MolmoAct2ClampNormalizedProcessorStep(),
|
||||
MolmoAct2PackInputsProcessorStep(
|
||||
checkpoint_path=config.checkpoint_path,
|
||||
checkpoint_revision=config.checkpoint_revision,
|
||||
checkpoint_force_download=config.checkpoint_force_download,
|
||||
trust_remote_code=config.trust_remote_code,
|
||||
action_mode=config.action_mode,
|
||||
discrete_action_tokenizer=config.discrete_action_tokenizer,
|
||||
image_keys=image_keys,
|
||||
setup_type=setup_type,
|
||||
control_mode=control_mode,
|
||||
normalize_language=config.normalize_language,
|
||||
add_setup_tokens=config.add_setup_tokens,
|
||||
add_control_tokens=config.add_control_tokens,
|
||||
num_state_tokens=config.num_state_tokens,
|
||||
max_sequence_length=config.max_sequence_length,
|
||||
chunk_size=chunk_size,
|
||||
max_action_dim=config.expected_max_action_dim,
|
||||
env_action_dim=env_action_dim,
|
||||
),
|
||||
DeviceProcessorStep(device=config.device),
|
||||
]
|
||||
|
||||
output_steps: list[ProcessorStep] = [
|
||||
MolmoAct2ClampActionProcessorStep(),
|
||||
MolmoAct2MaskedUnnormalizerProcessorStep(
|
||||
features=config.output_features,
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=masked_dataset_stats,
|
||||
),
|
||||
DeviceProcessorStep(device="cpu"),
|
||||
]
|
||||
|
||||
return (
|
||||
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
|
||||
steps=input_steps,
|
||||
name=POLICY_PREPROCESSOR_DEFAULT_NAME,
|
||||
),
|
||||
PolicyProcessorPipeline[PolicyAction, PolicyAction](
|
||||
steps=output_steps,
|
||||
name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
|
||||
to_transition=policy_action_to_transition,
|
||||
to_output=transition_to_policy_action,
|
||||
),
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -2915,6 +2915,10 @@ metaworld = [
|
||||
{ name = "scipy" },
|
||||
{ name = "torchcodec", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'AMD64' and sys_platform == 'linux') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'arm64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'win32'" },
|
||||
]
|
||||
molmoact2 = [
|
||||
{ name = "peft" },
|
||||
{ name = "transformers" },
|
||||
]
|
||||
motorbridge-dep = [
|
||||
{ name = "motorbridge" },
|
||||
]
|
||||
@@ -3128,6 +3132,7 @@ requires-dist = [
|
||||
{ name = "lerobot", extras = ["matplotlib-dep"], marker = "extra == 'sarm'" },
|
||||
{ name = "lerobot", extras = ["matplotlib-dep"], marker = "extra == 'unitree-g1'" },
|
||||
{ name = "lerobot", extras = ["metaworld"], marker = "extra == 'all'" },
|
||||
{ name = "lerobot", extras = ["molmoact2"], marker = "extra == 'all'" },
|
||||
{ name = "lerobot", extras = ["motorbridge-dep"], marker = "extra == 'rebot'" },
|
||||
{ name = "lerobot", extras = ["motorbridge-smart-servo-dep"], marker = "extra == 'rebot'" },
|
||||
{ name = "lerobot", extras = ["multi-task-dit"], marker = "extra == 'all'" },
|
||||
@@ -3135,6 +3140,7 @@ requires-dist = [
|
||||
{ name = "lerobot", extras = ["openarms"], marker = "extra == 'all'" },
|
||||
{ name = "lerobot", extras = ["peft"], marker = "extra == 'all'" },
|
||||
{ name = "lerobot", extras = ["peft-dep"], marker = "extra == 'groot'" },
|
||||
{ name = "lerobot", extras = ["peft-dep"], marker = "extra == 'molmoact2'" },
|
||||
{ name = "lerobot", extras = ["peft-dep"], marker = "extra == 'peft'" },
|
||||
{ name = "lerobot", extras = ["peft-dep"], marker = "extra == 'wallx'" },
|
||||
{ name = "lerobot", extras = ["phone"], marker = "extra == 'all'" },
|
||||
@@ -3172,6 +3178,7 @@ requires-dist = [
|
||||
{ name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'groot'" },
|
||||
{ name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'hilserl'" },
|
||||
{ name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'libero'" },
|
||||
{ name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'molmoact2'" },
|
||||
{ name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'multi-task-dit'" },
|
||||
{ name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'peft'" },
|
||||
{ name = "lerobot", extras = ["transformers-dep"], marker = "extra == 'pi'" },
|
||||
@@ -3244,7 +3251,7 @@ requires-dist = [
|
||||
{ name = "transformers", marker = "extra == 'transformers-dep'", specifier = ">=5.4.0,<5.6.0" },
|
||||
{ name = "wandb", marker = "extra == 'training'", specifier = ">=0.24.0,<0.25.0" },
|
||||
]
|
||||
provides-extras = ["dataset", "training", "hardware", "viz", "core-scripts", "evaluation", "dataset-viz", "av-dep", "pygame-dep", "placo-dep", "transformers-dep", "grpcio-dep", "can-dep", "peft-dep", "scipy-dep", "diffusers-dep", "qwen-vl-utils-dep", "matplotlib-dep", "pyserial-dep", "deepdiff-dep", "pynput-dep", "pyzmq-dep", "motorbridge-dep", "motorbridge-smart-servo-dep", "feetech", "dynamixel", "damiao", "robstride", "openarms", "gamepad", "hopejr", "lekiwi", "unitree-g1", "reachy2", "rebot", "kinematics", "intelrealsense", "phone", "diffusion", "wallx", "pi", "smolvla", "multi-task-dit", "groot", "sarm", "xvla", "eo1", "hilserl", "async", "peft", "dev", "notebook", "test", "video-benchmark", "aloha", "pusht", "libero", "metaworld", "all"]
|
||||
provides-extras = ["dataset", "training", "hardware", "viz", "core-scripts", "evaluation", "dataset-viz", "av-dep", "pygame-dep", "placo-dep", "transformers-dep", "grpcio-dep", "can-dep", "peft-dep", "scipy-dep", "diffusers-dep", "qwen-vl-utils-dep", "matplotlib-dep", "pyserial-dep", "deepdiff-dep", "pynput-dep", "pyzmq-dep", "motorbridge-dep", "motorbridge-smart-servo-dep", "feetech", "dynamixel", "damiao", "robstride", "openarms", "gamepad", "hopejr", "lekiwi", "unitree-g1", "reachy2", "rebot", "kinematics", "intelrealsense", "phone", "diffusion", "wallx", "pi", "molmoact2", "smolvla", "multi-task-dit", "groot", "sarm", "xvla", "eo1", "hilserl", "async", "peft", "dev", "notebook", "test", "video-benchmark", "aloha", "pusht", "libero", "metaworld", "all"]
|
||||
|
||||
[[package]]
|
||||
name = "librt"
|
||||
|
||||
Reference in New Issue
Block a user