mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-23 17:56:07 +00:00
feat(pi05): synthesize relative proprioceptive history
This commit is contained in:
Pepijn
@@ -0,0 +1,175 @@
|
||||
#!/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 numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
pytest.importorskip("transformers")
|
||||
|
||||
from lerobot.datasets.compute_stats import ( # noqa: E402
|
||||
compute_relative_action_stats,
|
||||
compute_state_history_stats,
|
||||
)
|
||||
from lerobot.policies.pi05.configuration_pi05 import PI05Config # noqa: E402
|
||||
from lerobot.policies.pi05.processor_pi05 import ( # noqa: E402
|
||||
Pi05FlattenStateHistoryProcessorStep,
|
||||
Pi05StateFromActionProcessorStep,
|
||||
)
|
||||
from lerobot.processor.relative_action_processor import ( # noqa: E402
|
||||
AbsoluteActionsProcessorStep,
|
||||
RelativeActionsProcessorStep,
|
||||
)
|
||||
from lerobot.types import TransitionKey # noqa: E402
|
||||
from lerobot.utils.constants import OBS_STATE # noqa: E402
|
||||
|
||||
|
||||
def _transition(action: torch.Tensor | None, state: torch.Tensor | None = None) -> dict:
|
||||
observation = {} if state is None else {OBS_STATE: state}
|
||||
return {
|
||||
TransitionKey.OBSERVATION: observation,
|
||||
TransitionKey.ACTION: action,
|
||||
TransitionKey.REWARD: None,
|
||||
TransitionKey.DONE: None,
|
||||
TransitionKey.TRUNCATED: None,
|
||||
TransitionKey.COMPLEMENTARY_DATA: {},
|
||||
}
|
||||
|
||||
|
||||
def test_pi05_config_requests_action_history_prefix():
|
||||
config = PI05Config(
|
||||
device="cpu",
|
||||
chunk_size=4,
|
||||
n_action_steps=4,
|
||||
state_from_action=True,
|
||||
proprioception_history_steps=2,
|
||||
)
|
||||
|
||||
assert config.action_delta_indices == [-1, 0, 1, 2, 3]
|
||||
|
||||
|
||||
def test_state_from_action_extracts_history_and_preserves_target_horizon():
|
||||
action = torch.arange(2 * 5 * 3, dtype=torch.float32).reshape(2, 5, 3)
|
||||
step = Pi05StateFromActionProcessorStep(enabled=True, history_steps=2)
|
||||
|
||||
result = step(_transition(action))
|
||||
|
||||
torch.testing.assert_close(result[TransitionKey.OBSERVATION][OBS_STATE], action[:, :2])
|
||||
torch.testing.assert_close(result[TransitionKey.ACTION], action[:, 1:])
|
||||
|
||||
|
||||
def test_relative_actions_use_newest_state_in_history_and_roundtrip():
|
||||
state_history = torch.tensor([[[1.0, 10.0], [2.0, 20.0]]])
|
||||
absolute = torch.tensor([[[3.0, 30.0], [4.0, 40.0]]])
|
||||
relative_step = RelativeActionsProcessorStep(enabled=True)
|
||||
absolute_step = AbsoluteActionsProcessorStep(enabled=True, relative_step=relative_step)
|
||||
|
||||
relative = relative_step(_transition(absolute, state_history))
|
||||
expected = torch.tensor([[[1.0, 10.0], [2.0, 20.0]]])
|
||||
torch.testing.assert_close(relative[TransitionKey.ACTION], expected)
|
||||
|
||||
recovered = absolute_step(_transition(relative[TransitionKey.ACTION]))
|
||||
torch.testing.assert_close(recovered[TransitionKey.ACTION], absolute)
|
||||
|
||||
|
||||
def test_flatten_state_history_preserves_chronological_order():
|
||||
state_history = torch.tensor([[[1.0, 2.0], [3.0, 4.0]]])
|
||||
step = Pi05FlattenStateHistoryProcessorStep(history_steps=2, max_state_dim=4)
|
||||
|
||||
result = step(_transition(torch.zeros(1, 2, 2), state_history))
|
||||
|
||||
torch.testing.assert_close(
|
||||
result[TransitionKey.OBSERVATION][OBS_STATE], torch.tensor([[1.0, 2.0, 3.0, 4.0]])
|
||||
)
|
||||
|
||||
|
||||
def test_state_history_can_be_relative_with_absolute_gripper():
|
||||
state_history = torch.tensor([[[1.0, 10.0, 0.2], [3.0, 20.0, 0.4]]])
|
||||
step = Pi05FlattenStateHistoryProcessorStep(
|
||||
history_steps=2,
|
||||
max_state_dim=6,
|
||||
relative=True,
|
||||
exclude_joints=["gripper"],
|
||||
state_names=["x", "y", "gripper_width"],
|
||||
)
|
||||
|
||||
result = step(_transition(torch.zeros(1, 2, 3), state_history))
|
||||
|
||||
torch.testing.assert_close(
|
||||
result[TransitionKey.OBSERVATION][OBS_STATE],
|
||||
torch.tensor([[-2.0, -10.0, 0.2, 0.0, 0.0, 0.4]]),
|
||||
)
|
||||
|
||||
|
||||
def test_inference_state_history_is_rolled_and_reset():
|
||||
step = Pi05StateFromActionProcessorStep(enabled=True, history_steps=2)
|
||||
|
||||
first = step(_transition(None, torch.tensor([[1.0, 2.0]])))
|
||||
second = step(_transition(None, torch.tensor([[3.0, 4.0]])))
|
||||
torch.testing.assert_close(
|
||||
first[TransitionKey.OBSERVATION][OBS_STATE], torch.tensor([[[1.0, 2.0], [1.0, 2.0]]])
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
second[TransitionKey.OBSERVATION][OBS_STATE], torch.tensor([[[1.0, 2.0], [3.0, 4.0]]])
|
||||
)
|
||||
|
||||
step.reset()
|
||||
reset = step(_transition(None, torch.tensor([[5.0, 6.0]])))
|
||||
torch.testing.assert_close(
|
||||
reset[TransitionKey.OBSERVATION][OBS_STATE], torch.tensor([[[5.0, 6.0], [5.0, 6.0]]])
|
||||
)
|
||||
|
||||
|
||||
def test_flatten_state_history_checks_max_state_dim():
|
||||
step = Pi05FlattenStateHistoryProcessorStep(history_steps=2, max_state_dim=3)
|
||||
|
||||
with pytest.raises(ValueError, match="above max_state_dim"):
|
||||
step(_transition(torch.zeros(1, 2, 2), torch.zeros(1, 2, 2)))
|
||||
|
||||
|
||||
def test_relative_stats_can_use_absolute_action_as_state():
|
||||
actions = np.asarray([[0.0, 0.0], [1.0, 2.0], [2.0, 4.0], [3.0, 6.0]], dtype=np.float32)
|
||||
dataset = {"action": actions, "episode_index": np.zeros(4, dtype=np.int64)}
|
||||
features = {"action": {"shape": [2], "names": ["x", "y"]}}
|
||||
|
||||
stats = compute_relative_action_stats(
|
||||
dataset,
|
||||
features,
|
||||
chunk_size=2,
|
||||
state_from_action=True,
|
||||
)
|
||||
|
||||
np.testing.assert_allclose(stats["mean"], [0.5, 1.0])
|
||||
|
||||
|
||||
def test_relative_state_history_stats_match_processor_representation():
|
||||
actions = np.asarray(
|
||||
[[0.0, 0.1], [1.0, 0.2], [3.0, 0.3]],
|
||||
dtype=np.float32,
|
||||
)
|
||||
dataset = {"action": actions, "episode_index": np.zeros(3, dtype=np.int64)}
|
||||
features = {"action": {"shape": [2], "names": ["x", "gripper_width"]}}
|
||||
|
||||
stats = compute_state_history_stats(
|
||||
dataset,
|
||||
features,
|
||||
history_steps=2,
|
||||
exclude_joints=["gripper"],
|
||||
relative=True,
|
||||
)
|
||||
|
||||
expected = np.asarray([[0.0, 0.1, 0.0, 0.1], [-1.0, 0.1, 0.0, 0.2], [-2.0, 0.2, 0.0, 0.3]])
|
||||
np.testing.assert_allclose(stats["mean"], expected.mean(axis=0))
|
||||
Reference in New Issue
Block a user