Merge remote-tracking branch 'hf/main' into feature/basic-peft-support

tests/utils/test_encoding_utils.py

@@ -1,6 +1,22 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import pytest
 
-from lerobot.utils.encoding_utils import (
+from lerobot.motors.encoding_utils import (
     decode_sign_magnitude,
     decode_twos_complement,
     encode_sign_magnitude,

tests/utils/test_io_utils.py

@@ -1,4 +1,6 @@
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -11,6 +13,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
 import json
 from pathlib import Path
 from typing import Any

tests/utils/test_logging_utils.py

@@ -1,4 +1,6 @@
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -11,6 +13,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
 import pytest
 
 from lerobot.utils.logging_utils import AverageMeter, MetricsTracker

tests/utils/test_process.py

@@ -22,7 +22,7 @@ from unittest.mock import patch
 
 import pytest
 
-from lerobot.utils.process import ProcessSignalHandler
+from lerobot.rl.process import ProcessSignalHandler
 
 
 # Fixture to reset shutdown_event_counter and original signal handlers before and after each test

tests/utils/test_queue.py

@@ -1,150 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import threading
-import time
-from queue import Queue
-
-from lerobot.utils.queue import get_last_item_from_queue
-
-
-def test_get_last_item_single_item():
-    """Test getting the last item when queue has only one item."""
-    queue = Queue()
-    queue.put("single_item")
-
-    result = get_last_item_from_queue(queue)
-
-    assert result == "single_item"
-    assert queue.empty()
-
-
-def test_get_last_item_multiple_items():
-    """Test getting the last item when queue has multiple items."""
-    queue = Queue()
-    items = ["first", "second", "third", "fourth", "last"]
-
-    for item in items:
-        queue.put(item)
-
-    result = get_last_item_from_queue(queue)
-
-    assert result == "last"
-    assert queue.empty()
-
-
-def test_get_last_item_different_types():
-    """Test with different data types in the queue."""
-    queue = Queue()
-    items = [1, 2.5, "string", {"key": "value"}, [1, 2, 3], ("tuple", "data")]
-
-    for item in items:
-        queue.put(item)
-
-    result = get_last_item_from_queue(queue)
-
-    assert result == ("tuple", "data")
-    assert queue.empty()
-
-
-def test_get_last_item_maxsize_queue():
-    """Test with a queue that has a maximum size."""
-    queue = Queue(maxsize=5)
-
-    # Fill the queue
-    for i in range(5):
-        queue.put(i)
-
-    # Give the queue time to fill
-    time.sleep(0.1)
-
-    result = get_last_item_from_queue(queue)
-
-    assert result == 4
-    assert queue.empty()
-
-
-def test_get_last_item_with_none_values():
-    """Test with None values in the queue."""
-    queue = Queue()
-    items = [1, None, 2, None, 3]
-
-    for item in items:
-        queue.put(item)
-
-    # Give the queue time to fill
-    time.sleep(0.1)
-
-    result = get_last_item_from_queue(queue)
-
-    assert result == 3
-    assert queue.empty()
-
-
-def test_get_last_item_blocking_timeout():
-    """Test get_last_item_from_queue returns None on timeout."""
-    queue = Queue()
-    result = get_last_item_from_queue(queue, block=True, timeout=0.1)
-    assert result is None
-
-
-def test_get_last_item_non_blocking_empty():
-    """Test get_last_item_from_queue with block=False on an empty queue returns None."""
-    queue = Queue()
-    result = get_last_item_from_queue(queue, block=False)
-    assert result is None
-
-
-def test_get_last_item_non_blocking_success():
-    """Test get_last_item_from_queue with block=False on a non-empty queue."""
-    queue = Queue()
-    items = ["first", "second", "last"]
-    for item in items:
-        queue.put(item)
-
-    # Give the queue time to fill
-    time.sleep(0.1)
-
-    result = get_last_item_from_queue(queue, block=False)
-    assert result == "last"
-    assert queue.empty()
-
-
-def test_get_last_item_blocking_waits_for_item():
-    """Test that get_last_item_from_queue waits for an item if block=True."""
-    queue = Queue()
-    result = []
-
-    def producer():
-        queue.put("item1")
-        queue.put("item2")
-
-    def consumer():
-        # This will block until the producer puts the first item
-        item = get_last_item_from_queue(queue, block=True, timeout=0.2)
-        result.append(item)
-
-    producer_thread = threading.Thread(target=producer)
-    consumer_thread = threading.Thread(target=consumer)
-
-    producer_thread.start()
-    consumer_thread.start()
-
-    producer_thread.join()
-    consumer_thread.join()
-
-    assert result == ["item2"]
-    assert queue.empty()

tests/utils/test_random_utils.py

@@ -1,4 +1,6 @@
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -11,6 +13,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
 import random
 
 import numpy as np

tests/utils/test_replay_buffer.py

@@ -15,18 +15,19 @@
 # limitations under the License.
 
 import sys
-from typing import Callable
+from collections.abc import Callable
 
 import pytest
 import torch
 
 from lerobot.datasets.lerobot_dataset import LeRobotDataset
-from lerobot.utils.buffer import BatchTransition, ReplayBuffer, random_crop_vectorized
+from lerobot.rl.buffer import BatchTransition, ReplayBuffer, random_crop_vectorized
+from lerobot.utils.constants import ACTION, DONE, OBS_IMAGE, OBS_STATE, OBS_STR, REWARD
 from tests.fixtures.constants import DUMMY_REPO_ID
 
 
 def state_dims() -> list[str]:
-    return ["observation.image", "observation.state"]
+    return [OBS_IMAGE, OBS_STATE]
 
 
 @pytest.fixture
@@ -61,10 +62,10 @@ def create_random_image() -> torch.Tensor:
 
 def create_dummy_transition() -> dict:
     return {
-        "observation.image": create_random_image(),
-        "action": torch.randn(4),
+        OBS_IMAGE: create_random_image(),
+        ACTION: torch.randn(4),
         "reward": torch.tensor(1.0),
-        "observation.state": torch.randn(
+        OBS_STATE: torch.randn(
             10,
         ),
         "done": torch.tensor(False),
@@ -98,8 +99,8 @@ def create_dataset_from_replay_buffer(tmp_path) -> tuple[LeRobotDataset, ReplayB
 
 def create_dummy_state() -> dict:
     return {
-        "observation.image": create_random_image(),
-        "observation.state": torch.randn(
+        OBS_IMAGE: create_random_image(),
+        OBS_STATE: torch.randn(
             10,
         ),
     }
@@ -120,7 +121,7 @@ def get_tensors_memory_consumption(obj, visited_addresses):
 
     if isinstance(obj, torch.Tensor):
         return get_tensor_memory_consumption(obj)
-    elif isinstance(obj, (list, tuple)):
+    elif isinstance(obj, (list | tuple)):
         for item in obj:
             total_size += get_tensors_memory_consumption(item, visited_addresses)
     elif isinstance(obj, dict):
@@ -180,7 +181,7 @@ def test_empty_buffer_sample_raises_error(replay_buffer):
 
 def test_zero_capacity_buffer_raises_error():
     with pytest.raises(ValueError, match="Capacity must be greater than 0."):
-        ReplayBuffer(0, "cpu", ["observation", "next_observation"])
+        ReplayBuffer(0, "cpu", [OBS_STR, "next_observation"])
 
 
 def test_add_transition(replay_buffer, dummy_state, dummy_action):
@@ -203,7 +204,7 @@ def test_add_transition(replay_buffer, dummy_state, dummy_action):
 
 
 def test_add_over_capacity():
-    replay_buffer = ReplayBuffer(2, "cpu", ["observation", "next_observation"])
+    replay_buffer = ReplayBuffer(2, "cpu", [OBS_STR, "next_observation"])
     dummy_state_1 = create_dummy_state()
     dummy_action_1 = create_dummy_action()
 
@@ -340,7 +341,7 @@ def test_sample_batch(replay_buffer):
                     f"{k} should be equal to one of the dummy states."
                 )
 
-    for got_action_item in got_batch_transition["action"]:
+    for got_action_item in got_batch_transition[ACTION]:
         assert any(torch.equal(got_action_item, dummy_action) for dummy_action in dummy_actions), (
             "Actions should be equal to the dummy actions."
         )
@@ -373,22 +374,22 @@ def test_to_lerobot_dataset(tmp_path):
     assert ds.num_frames == 4
 
     for j, value in enumerate(ds):
-        print(torch.equal(value["observation.image"], buffer.next_states["observation.image"][j]))
+        print(torch.equal(value[OBS_IMAGE], buffer.next_states[OBS_IMAGE][j]))
 
     for i in range(len(ds)):
         for feature, value in ds[i].items():
-            if feature == "action":
+            if feature == ACTION:
                 assert torch.equal(value, buffer.actions[i])
-            elif feature == "next.reward":
+            elif feature == REWARD:
                 assert torch.equal(value, buffer.rewards[i])
-            elif feature == "next.done":
+            elif feature == DONE:
                 assert torch.equal(value, buffer.dones[i])
-            elif feature == "observation.image":
-                # Tenssor -> numpy is not precise, so we have some diff there
+            elif feature == OBS_IMAGE:
+                # Tensor -> numpy is not precise, so we have some diff there
                 # TODO: Check and fix it
-                torch.testing.assert_close(value, buffer.states["observation.image"][i], rtol=0.3, atol=0.003)
-            elif feature == "observation.state":
-                assert torch.equal(value, buffer.states["observation.state"][i])
+                torch.testing.assert_close(value, buffer.states[OBS_IMAGE][i], rtol=0.3, atol=0.003)
+            elif feature == OBS_STATE:
+                assert torch.equal(value, buffer.states[OBS_STATE][i])
 
 
 def test_from_lerobot_dataset(tmp_path):
@@ -436,14 +437,14 @@ def test_from_lerobot_dataset(tmp_path):
     )
 
     assert torch.equal(
-        replay_buffer.states["observation.state"][: len(replay_buffer)],
-        reconverted_buffer.states["observation.state"][: len(replay_buffer)],
+        replay_buffer.states[OBS_STATE][: len(replay_buffer)],
+        reconverted_buffer.states[OBS_STATE][: len(replay_buffer)],
     ), "State should be the same after converting to dataset and return back"
 
     for i in range(4):
         torch.testing.assert_close(
-            replay_buffer.states["observation.image"][i],
-            reconverted_buffer.states["observation.image"][i],
+            replay_buffer.states[OBS_IMAGE][i],
+            reconverted_buffer.states[OBS_IMAGE][i],
             rtol=0.4,
             atol=0.004,
         )
@@ -454,16 +455,16 @@ def test_from_lerobot_dataset(tmp_path):
         next_index = (i + 1) % 4
 
         torch.testing.assert_close(
-            replay_buffer.states["observation.image"][next_index],
-            reconverted_buffer.next_states["observation.image"][i],
+            replay_buffer.states[OBS_IMAGE][next_index],
+            reconverted_buffer.next_states[OBS_IMAGE][i],
             rtol=0.4,
             atol=0.004,
         )
 
     for i in range(2, 4):
         assert torch.equal(
-            replay_buffer.states["observation.state"][i],
-            reconverted_buffer.next_states["observation.state"][i],
+            replay_buffer.states[OBS_STATE][i],
+            reconverted_buffer.next_states[OBS_STATE][i],
         )
 
 
@@ -494,7 +495,7 @@ def test_buffer_sample_alignment():
 
     for i in range(50):
         state_sig = batch["state"]["state_value"][i].item()
-        action_val = batch["action"][i].item()
+        action_val = batch[ACTION][i].item()
         reward_val = batch["reward"][i].item()
         next_state_sig = batch["next_state"]["state_value"][i].item()
         is_done = batch["done"][i].item() > 0.5
@@ -563,10 +564,8 @@ def test_check_image_augmentations_with_drq_and_dummy_image_augmentation_functio
     replay_buffer.add(dummy_state, dummy_action, 1.0, dummy_state, False, False)
 
     sampled_transitions = replay_buffer.sample(1)
-    assert torch.all(sampled_transitions["state"]["observation.image"] == 10), (
-        "Image augmentations should be applied"
-    )
-    assert torch.all(sampled_transitions["next_state"]["observation.image"] == 10), (
+    assert torch.all(sampled_transitions["state"][OBS_IMAGE] == 10), "Image augmentations should be applied"
+    assert torch.all(sampled_transitions["next_state"][OBS_IMAGE] == 10), (
         "Image augmentations should be applied"
     )
 
@@ -580,8 +579,8 @@ def test_check_image_augmentations_with_drq_and_default_image_augmentation_funct
 
     # Let's check that it doesn't fail and shapes are correct
     sampled_transitions = replay_buffer.sample(1)
-    assert sampled_transitions["state"]["observation.image"].shape == (1, 3, 84, 84)
-    assert sampled_transitions["next_state"]["observation.image"].shape == (1, 3, 84, 84)
+    assert sampled_transitions["state"][OBS_IMAGE].shape == (1, 3, 84, 84)
+    assert sampled_transitions["next_state"][OBS_IMAGE].shape == (1, 3, 84, 84)
 
 
 def test_random_crop_vectorized_basic():
@@ -620,7 +619,7 @@ def _populate_buffer_for_async_test(capacity: int = 10) -> ReplayBuffer:
     buffer = ReplayBuffer(
         capacity=capacity,
         device="cpu",
-        state_keys=["observation.image", "observation.state"],
+        state_keys=[OBS_IMAGE, OBS_STATE],
         storage_device="cpu",
     )
 
@@ -628,8 +627,8 @@ def _populate_buffer_for_async_test(capacity: int = 10) -> ReplayBuffer:
         img = torch.ones(3, 128, 128) * i
         state_vec = torch.arange(11).float() + i
         state = {
-            "observation.image": img,
-            "observation.state": state_vec,
+            OBS_IMAGE: img,
+            OBS_STATE: state_vec,
         }
         buffer.add(
             state=state,
@@ -648,14 +647,14 @@ def test_async_iterator_shapes_basic():
     iterator = buffer.get_iterator(batch_size=batch_size, async_prefetch=True, queue_size=1)
     batch = next(iterator)
 
-    images = batch["state"]["observation.image"]
-    states = batch["state"]["observation.state"]
+    images = batch["state"][OBS_IMAGE]
+    states = batch["state"][OBS_STATE]
 
     assert images.shape == (batch_size, 3, 128, 128)
     assert states.shape == (batch_size, 11)
 
-    next_images = batch["next_state"]["observation.image"]
-    next_states = batch["next_state"]["observation.state"]
+    next_images = batch["next_state"][OBS_IMAGE]
+    next_states = batch["next_state"][OBS_STATE]
 
     assert next_images.shape == (batch_size, 3, 128, 128)
     assert next_states.shape == (batch_size, 11)
@@ -668,13 +667,13 @@ def test_async_iterator_multiple_iterations():
 
     for _ in range(5):
         batch = next(iterator)
-        images = batch["state"]["observation.image"]
-        states = batch["state"]["observation.state"]
+        images = batch["state"][OBS_IMAGE]
+        states = batch["state"][OBS_STATE]
         assert images.shape == (batch_size, 3, 128, 128)
         assert states.shape == (batch_size, 11)
 
-        next_images = batch["next_state"]["observation.image"]
-        next_states = batch["next_state"]["observation.state"]
+        next_images = batch["next_state"][OBS_IMAGE]
+        next_states = batch["next_state"][OBS_STATE]
         assert next_images.shape == (batch_size, 3, 128, 128)
         assert next_states.shape == (batch_size, 11)
 

tests/utils/test_train_utils.py

@@ -1,4 +1,6 @@
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -11,10 +13,11 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
 from pathlib import Path
 from unittest.mock import Mock, patch
 
-from lerobot.constants import (
+from lerobot.utils.constants import (
     CHECKPOINTS_DIR,
     LAST_CHECKPOINT_LINK,
     OPTIMIZER_PARAM_GROUPS,

tests/utils/test_visualization_utils.py

@@ -0,0 +1,226 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import importlib
+import sys
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+
+from lerobot.processor import TransitionKey
+from lerobot.utils.constants import OBS_STATE
+
+
+@pytest.fixture
+def mock_rerun(monkeypatch):
+    """
+    Provide a mock `rerun` module so tests don't depend on the real library.
+    Also reload the module-under-test so it binds to this mock `rr`.
+    """
+    calls = []
+
+    class DummyScalar:
+        def __init__(self, value):
+            self.value = float(value)
+
+    class DummyImage:
+        def __init__(self, arr):
+            self.arr = arr
+
+    def dummy_log(key, obj, **kwargs):
+        calls.append((key, obj, kwargs))
+
+    dummy_rr = SimpleNamespace(
+        Scalar=DummyScalar,
+        Image=DummyImage,
+        log=dummy_log,
+        init=lambda *a, **k: None,
+        spawn=lambda *a, **k: None,
+    )
+
+    # Inject fake module into sys.modules
+    monkeypatch.setitem(sys.modules, "rerun", dummy_rr)
+
+    # Now import and reload the module under test, to bind to our rerun mock
+    import lerobot.utils.visualization_utils as vu
+
+    importlib.reload(vu)
+
+    # Expose both the reloaded module and the call recorder
+    yield vu, calls
+
+
+def _keys(calls):
+    """Helper to extract just the keys logged to rr.log"""
+    return [k for (k, _obj, _kw) in calls]
+
+
+def _obj_for(calls, key):
+    """Find the first object logged under a given key."""
+    for k, obj, _kw in calls:
+        if k == key:
+            return obj
+    raise KeyError(f"Key {key} not found in calls: {calls}")
+
+
+def _kwargs_for(calls, key):
+    for k, _obj, kw in calls:
+        if k == key:
+            return kw
+    raise KeyError(f"Key {key} not found in calls: {calls}")
+
+
+def test_log_rerun_data_envtransition_scalars_and_image(mock_rerun):
+    vu, calls = mock_rerun
+
+    # Build EnvTransition dict
+    obs = {
+        f"{OBS_STATE}.temperature": np.float32(25.0),
+        # CHW image should be converted to HWC for rr.Image
+        "observation.camera": np.zeros((3, 10, 20), dtype=np.uint8),
+    }
+    act = {
+        "action.throttle": 0.7,
+        # 1D array should log individual Scalars with suffix _i
+        "action.vector": np.array([1.0, 2.0], dtype=np.float32),
+    }
+    transition = {
+        TransitionKey.OBSERVATION: obs,
+        TransitionKey.ACTION: act,
+    }
+
+    # Extract observation and action data from transition like in the real call sites
+    obs_data = transition.get(TransitionKey.OBSERVATION, {})
+    action_data = transition.get(TransitionKey.ACTION, {})
+    vu.log_rerun_data(observation=obs_data, action=action_data)
+
+    # We expect:
+    # - observation.state.temperature -> Scalar
+    # - observation.camera -> Image (HWC) with static=True
+    # - action.throttle -> Scalar
+    # - action.vector_0, action.vector_1 -> Scalars
+    expected_keys = {
+        f"{OBS_STATE}.temperature",
+        "observation.camera",
+        "action.throttle",
+        "action.vector_0",
+        "action.vector_1",
+    }
+    assert set(_keys(calls)) == expected_keys
+
+    # Check scalar types and values
+    temp_obj = _obj_for(calls, f"{OBS_STATE}.temperature")
+    assert type(temp_obj).__name__ == "DummyScalar"
+    assert temp_obj.value == pytest.approx(25.0)
+
+    throttle_obj = _obj_for(calls, "action.throttle")
+    assert type(throttle_obj).__name__ == "DummyScalar"
+    assert throttle_obj.value == pytest.approx(0.7)
+
+    v0 = _obj_for(calls, "action.vector_0")
+    v1 = _obj_for(calls, "action.vector_1")
+    assert type(v0).__name__ == "DummyScalar"
+    assert type(v1).__name__ == "DummyScalar"
+    assert v0.value == pytest.approx(1.0)
+    assert v1.value == pytest.approx(2.0)
+
+    # Check image handling: CHW -> HWC
+    img_obj = _obj_for(calls, "observation.camera")
+    assert type(img_obj).__name__ == "DummyImage"
+    assert img_obj.arr.shape == (10, 20, 3)  # transposed
+    assert _kwargs_for(calls, "observation.camera").get("static", False) is True  # static=True for images
+
+
+def test_log_rerun_data_plain_list_ordering_and_prefixes(mock_rerun):
+    vu, calls = mock_rerun
+
+    # First dict without prefixes treated as observation
+    # Second dict without prefixes treated as action
+    obs_plain = {
+        "temp": 1.5,
+        # Already HWC image => should stay as-is
+        "img": np.zeros((5, 6, 3), dtype=np.uint8),
+        "none": None,  # should be skipped
+    }
+    act_plain = {
+        "throttle": 0.3,
+        "vec": np.array([9, 8, 7], dtype=np.float32),
+    }
+
+    # Extract observation and action data from list like the old function logic did
+    # First dict was treated as observation, second as action
+    vu.log_rerun_data(observation=obs_plain, action=act_plain)
+
+    # Expected keys with auto-prefixes
+    expected = {
+        "observation.temp",
+        "observation.img",
+        "action.throttle",
+        "action.vec_0",
+        "action.vec_1",
+        "action.vec_2",
+    }
+    logged = set(_keys(calls))
+    assert logged == expected
+
+    # Scalars
+    t = _obj_for(calls, "observation.temp")
+    assert type(t).__name__ == "DummyScalar"
+    assert t.value == pytest.approx(1.5)
+
+    throttle = _obj_for(calls, "action.throttle")
+    assert type(throttle).__name__ == "DummyScalar"
+    assert throttle.value == pytest.approx(0.3)
+
+    # Image stays HWC
+    img = _obj_for(calls, "observation.img")
+    assert type(img).__name__ == "DummyImage"
+    assert img.arr.shape == (5, 6, 3)
+    assert _kwargs_for(calls, "observation.img").get("static", False) is True
+
+    # Vectors
+    for i, val in enumerate([9, 8, 7]):
+        o = _obj_for(calls, f"action.vec_{i}")
+        assert type(o).__name__ == "DummyScalar"
+        assert o.value == pytest.approx(val)
+
+
+def test_log_rerun_data_kwargs_only(mock_rerun):
+    vu, calls = mock_rerun
+
+    vu.log_rerun_data(
+        observation={"observation.temp": 10.0, "observation.gray": np.zeros((8, 8, 1), dtype=np.uint8)},
+        action={"action.a": 1.0},
+    )
+
+    keys = set(_keys(calls))
+    assert "observation.temp" in keys
+    assert "observation.gray" in keys
+    assert "action.a" in keys
+
+    temp = _obj_for(calls, "observation.temp")
+    assert type(temp).__name__ == "DummyScalar"
+    assert temp.value == pytest.approx(10.0)
+
+    img = _obj_for(calls, "observation.gray")
+    assert type(img).__name__ == "DummyImage"
+    assert img.arr.shape == (8, 8, 1)  # remains HWC
+    assert _kwargs_for(calls, "observation.gray").get("static", False) is True
+
+    a = _obj_for(calls, "action.a")
+    assert type(a).__name__ == "DummyScalar"
+    assert a.value == pytest.approx(1.0)