fix rollout policy revision loading

Co-authored-by: RaviTeja-Kondeti <rkondet3@asu.edu>

tests/test_rollout.py

@@ -17,6 +17,8 @@
 from __future__ import annotations
 
 import dataclasses
+import sys
+from types import SimpleNamespace
 from unittest.mock import MagicMock
 
 import pytest
@@ -106,6 +108,109 @@ def test_sentry_config_defaults():
     assert cfg.target_video_file_size_mb is None
 
 
+def test_rollout_config_passes_policy_pretrained_revision(monkeypatch):
+    from lerobot.configs import PreTrainedConfig, parser
+    from lerobot.rollout import RolloutConfig
+    from tests.mocks.mock_robot import MockRobotConfig
+
+    captured = {}
+
+    def fake_from_pretrained(cls, pretrained_name_or_path, **kwargs):
+        captured["pretrained_name_or_path"] = pretrained_name_or_path
+        captured.update(kwargs)
+        return SimpleNamespace(device="cpu", pretrained_revision=kwargs["revision"])
+
+    monkeypatch.setattr(parser, "get_yaml_overrides", lambda _: ["--pretrained_revision=yaml-sha"])
+    monkeypatch.setattr(
+        sys,
+        "argv",
+        ["lerobot-rollout", "--policy.path=user/policy", "--policy.pretrained_revision=cli-sha"],
+    )
+    monkeypatch.setattr(PreTrainedConfig, "from_pretrained", classmethod(fake_from_pretrained))
+
+    cfg = RolloutConfig(robot=MockRobotConfig())
+
+    assert captured["pretrained_name_or_path"] == "user/policy"
+    assert captured["revision"] == "cli-sha"
+    assert captured["cli_overrides"] == [
+        "--pretrained_revision=yaml-sha",
+        "--pretrained_revision=cli-sha",
+    ]
+    assert cfg.policy.pretrained_path == "user/policy"
+    assert cfg.policy.pretrained_revision == "cli-sha"
+
+
+def test_load_pretrained_policy_passes_revision(monkeypatch):
+    import lerobot.rollout.context as rollout_context
+
+    policy_config = SimpleNamespace(
+        type="mock",
+        use_peft=False,
+        pretrained_path="user/policy",
+        pretrained_revision="policy-sha",
+    )
+    policy_class = MagicMock()
+    loaded_policy = MagicMock()
+    policy_class.from_pretrained.return_value = loaded_policy
+    monkeypatch.setattr(rollout_context, "get_policy_class", lambda _: policy_class)
+
+    policy = rollout_context._load_pretrained_policy(policy_config)
+
+    assert policy is loaded_policy
+    policy_class.from_pretrained.assert_called_once_with(
+        "user/policy",
+        config=policy_config,
+        revision="policy-sha",
+    )
+
+
+def test_load_pretrained_peft_policy_keeps_adapter_and_base_revisions_separate(monkeypatch):
+    import lerobot.rollout.context as rollout_context
+
+    policy_config = SimpleNamespace(
+        type="mock",
+        use_peft=True,
+        pretrained_path="user/adapter",
+        pretrained_revision="adapter-sha",
+    )
+    policy_class = MagicMock()
+    base_policy = MagicMock()
+    policy_class.from_pretrained.return_value = base_policy
+    monkeypatch.setattr(rollout_context, "get_policy_class", lambda _: policy_class)
+
+    peft_config = SimpleNamespace(
+        base_model_name_or_path="user/base-policy",
+        revision="base-sha",
+    )
+    peft_config_from_pretrained = MagicMock(return_value=peft_config)
+    adapted_policy = MagicMock()
+    peft_model_from_pretrained = MagicMock(return_value=adapted_policy)
+    monkeypatch.setitem(
+        sys.modules,
+        "peft",
+        SimpleNamespace(
+            PeftConfig=SimpleNamespace(from_pretrained=peft_config_from_pretrained),
+            PeftModel=SimpleNamespace(from_pretrained=peft_model_from_pretrained),
+        ),
+    )
+
+    policy = rollout_context._load_pretrained_policy(policy_config)
+
+    assert policy is adapted_policy
+    peft_config_from_pretrained.assert_called_once_with("user/adapter", revision="adapter-sha")
+    policy_class.from_pretrained.assert_called_once_with(
+        pretrained_name_or_path="user/base-policy",
+        config=policy_config,
+        revision="base-sha",
+    )
+    peft_model_from_pretrained.assert_called_once_with(
+        base_policy,
+        "user/adapter",
+        config=peft_config,
+        revision="adapter-sha",
+    )
+
+
 # ---------------------------------------------------------------------------
 # RolloutRingBuffer
 # ---------------------------------------------------------------------------