admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-15 14:02:14 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor(policies): clean MolmoAct2 to follow EO1/TOPReward patterns

Browse Source 
Align the MolmoAct2 implementation with lerobot codebase conventions:

- Rename hf_model/ to molmoact2_hf_model/
- Slim config: move all I/O and runtime logic to modeling
- Remove blanket  from 8 vendored files, fix 66 lint issues
- Deduplicate _hf_token() and _resolve_checkpoint_location()
- Make huggingface_hub imports lazy
- Remove custom MolmoAct2CosineDecayWithWarmupSchedulerConfig, use base class
- Extract 13 static/classmethods from MolmoAct2Policy to free functions
- Replace print() with logger in vendored action_tokenizer
- Add module docstrings, class docstring, and key method docstrings
- Add module-level loggers to modeling and processor
- Fix docs: pip to uv install, deduplicate README symlink
- Remove shebangs from all files
This commit is contained in:
Khalil Meftah
2026-06-05 16:31:03 +02:00
parent 2e9cd87bbd
commit 17e217d175
15 changed files with 611 additions and 694 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/policies/molmoact2/test_molmoact2.py
+35 -48
View File
@@ -1,5 +1,3 @@
#!/usr/bin/env python
# Copyright 2026 The Allen Institute for Artificial Intelligence and The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
@@ -35,16 +33,16 @@ pytest.importorskip("scipy")
from lerobot.configs import FeatureType, NormalizationMode, PolicyFeature
from lerobot.policies import get_policy_class, make_policy_config
from lerobot.policies.molmoact2 import (
configuration_molmoact2 as molmoact2_config,
modeling_molmoact2 as molmoact2_modeling,
processor_molmoact2 as molmoact2_processor,
)
from lerobot.policies.molmoact2.configuration_molmoact2 import (
MolmoAct2Config,
MolmoAct2CosineDecayWithWarmupSchedulerConfig,
infer_molmoact2_max_sequence_length,
from lerobot.policies.molmoact2.configuration_molmoact2 import MolmoAct2Config
from lerobot.policies.molmoact2.modeling_molmoact2 import (
MolmoAct2Policy,
_apply_action_chunk_padding_mask,
_apply_action_dim_padding_mask,
_combine_rollout_seeds,
)
from lerobot.policies.molmoact2.modeling_molmoact2 import MolmoAct2Policy
from lerobot.policies.molmoact2.processor_molmoact2 import (
MolmoAct2ClampNormalizedProcessorStep,
MolmoAct2MaskedNormalizerProcessorStep,
@@ -53,6 +51,7 @@ from lerobot.policies.molmoact2.processor_molmoact2 import (
_add_gripper_masks_to_stats,
_build_discrete_state_string,
_normalize_question_text,
infer_molmoact2_max_sequence_length,
make_molmoact2_pre_post_processors,
)
from lerobot.policies.rtc.configuration_rtc import RTCConfig
@@ -71,34 +70,38 @@ def test_molmoact2_policy_registration():
assert cfg.per_episode_seed is False
assert cfg.eval_seed is None
assert cfg.normalize_language is True
assert cfg.get_scheduler_preset().num_decay_steps is None
assert cfg.get_scheduler_preset().num_decay_steps == 100_000
assert cfg.action_delta_indices == list(range(cfg.chunk_size))
assert get_policy_class("molmoact2") is MolmoAct2Policy
def test_molmoact2_checkpoint_download_ignores_remote_python(monkeypatch):
import huggingface_hub
download_kwargs = {}
def fake_snapshot_download(**kwargs):
download_kwargs.update(kwargs)
return "/tmp/downloaded-molmoact2"
monkeypatch.setattr(molmoact2_config, "snapshot_download", fake_snapshot_download)
monkeypatch.setattr(huggingface_hub, "snapshot_download", fake_snapshot_download)
checkpoint_location = molmoact2_config._resolve_checkpoint_location("allenai/MolmoAct2")
checkpoint_location = molmoact2_modeling._resolve_checkpoint_location("allenai/MolmoAct2")
assert checkpoint_location == "/tmp/downloaded-molmoact2"
assert download_kwargs["ignore_patterns"] == ["*.py", "*.pyc", "__pycache__/*"]
def test_molmoact2_scheduler_decay_steps_auto_match_training_steps():
def test_molmoact2_scheduler_auto_scales_to_training_steps():
from lerobot.optim import CosineDecayWithWarmupSchedulerConfig
param = torch.nn.Parameter(torch.ones(()))
optimizer = torch.optim.AdamW([param], lr=0.001)
config = MolmoAct2CosineDecayWithWarmupSchedulerConfig(
config = CosineDecayWithWarmupSchedulerConfig(
peak_lr=0.01,
decay_lr=0.001,
num_warmup_steps=10,
num_decay_steps=None,
num_decay_steps=100_000,
)
scheduler = config.build(optimizer, num_training_steps=100)
@@ -123,9 +126,7 @@ def test_molmoact2_rollout_generator_uses_eval_seed_per_task():
batch_size=3,
device=torch.device("cpu"),
)
expected_first = torch.Generator().manual_seed(
MolmoAct2Policy._combine_rollout_seeds(first_seed=1000, batch_size=3)
)
expected_first = torch.Generator().manual_seed(_combine_rollout_seeds(first_seed=1000, batch_size=3))
assert torch.allclose(torch.rand(4, generator=first), torch.rand(4, generator=expected_first))
policy.reset()
@@ -134,9 +135,7 @@ def test_molmoact2_rollout_generator_uses_eval_seed_per_task():
batch_size=3,
device=torch.device("cpu"),
)
expected_second = torch.Generator().manual_seed(
MolmoAct2Policy._combine_rollout_seeds(first_seed=1003, batch_size=3)
)
expected_second = torch.Generator().manual_seed(_combine_rollout_seeds(first_seed=1003, batch_size=3))
assert torch.allclose(torch.rand(4, generator=second), torch.rand(4, generator=expected_second))
policy.reset()
@@ -145,9 +144,7 @@ def test_molmoact2_rollout_generator_uses_eval_seed_per_task():
batch_size=3,
device=torch.device("cpu"),
)
expected_new_task = torch.Generator().manual_seed(
MolmoAct2Policy._combine_rollout_seeds(first_seed=1000, batch_size=3)
)
expected_new_task = torch.Generator().manual_seed(_combine_rollout_seeds(first_seed=1000, batch_size=3))
assert torch.allclose(torch.rand(4, generator=new_task), torch.rand(4, generator=expected_new_task))
@@ -537,36 +534,26 @@ def test_train_action_expert_only_requires_continuous_action_mode():
def test_molmoact2_sequence_length_is_inferred_from_fixed_token_budget():
cfg = MolmoAct2Config(
action_mode="both",
chunk_size=10,
n_action_steps=10,
image_keys=["observation.images.image", "observation.images.wrist_image"],
input_features={OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(8,))},
output_features={ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(7,))},
)
assert cfg.max_sequence_length is None
assert cfg.inferred_max_sequence_length() == 640
assert cfg.inferred_max_sequence_length(include_discrete_action=False) == 576
assert (
infer_molmoact2_max_sequence_length(
num_images=2,
state_dim=8,
action_dim=7,
action_horizon=30,
include_discrete_action=True,
num_images=2, state_dim=8, action_dim=7, action_horizon=10, include_discrete_action=True
)
== 640
)
assert (
infer_molmoact2_max_sequence_length(
num_images=2, state_dim=8, action_dim=7, action_horizon=10, include_discrete_action=False
)
== 576
)
assert (
infer_molmoact2_max_sequence_length(
num_images=2, state_dim=8, action_dim=7, action_horizon=30, include_discrete_action=True
)
== 768
)
def test_molmoact2_sequence_length_override_is_preserved():
cfg = MolmoAct2Config(max_sequence_length=1024)
assert cfg.inferred_max_sequence_length(num_images=2, state_dim=8, action_dim=7) == 1024
def test_train_action_expert_only_freezes_non_action_expert_params():
class DummyBackbone(torch.nn.Module):
def __init__(self):
@@ -963,7 +950,7 @@ def test_action_dim_padding_loss_reduces_like_old_trainer():
]
)
reduced = MolmoAct2Policy._apply_action_dim_padding_mask(loss, action_dim_is_pad)
reduced = _apply_action_dim_padding_mask(loss, action_dim_is_pad)
expected = torch.stack(
[
@@ -979,7 +966,7 @@ def test_action_chunk_padding_keeps_old_mean_denominator():
loss = torch.ones(1, 2, 4, 3)
action_horizon_is_pad = torch.tensor([[False, False, True, True]])
masked = MolmoAct2Policy._apply_action_chunk_padding_mask(loss, action_horizon_is_pad)
masked = _apply_action_chunk_padding_mask(loss, action_horizon_is_pad)
assert masked.mean().item() == 0.5