diff --git a/tests/policies/test_inference_state_precision.py b/tests/policies/test_inference_state_precision.py index 3b4364c3..e3a2545e 100644 --- a/tests/policies/test_inference_state_precision.py +++ b/tests/policies/test_inference_state_precision.py @@ -46,7 +46,7 @@ from opentau.configs.types import FeatureType, NormalizationMode, PolicyFeature from opentau.datasets.action_indexing import subtract_chunk_start_state from opentau.policies.layers import PerGroupLinear -from opentau.policies.normalize import Normalize +from opentau.policies.normalize import Normalize, Unnormalize from opentau.policies.pi05 import modeling_pi05 from opentau.policies.pi05.configuration_pi05 import PI05Config from opentau.policies.pi05.modeling_pi05 import PI05Policy @@ -55,6 +55,7 @@ from opentau.scripts.grpc.server import RobotPolicyServicer from opentau.scripts.robocasa import server as robocasa_server from opentau.utils.utils import INFERENCE_STATE_DTYPE, create_dummy_observation +from tests.utils import assert_action_normalization_round_trips STATE_DIM = 8 # 7 arm joints + gripper; already max_state_dim, so the servers' pad is a no-op ACTION_DIM = 8 @@ -373,6 +374,41 @@ def test_normalize_computes_in_its_stats_dtype(mode, stats): assert torch.equal(upcast, norm32({"state": x.to(torch.bfloat16).float()}, index)["state"]) +@pytest.mark.parametrize( + "dtype, mean_error", + [(torch.bfloat16, 1 / 16), (torch.float32, 1e-3)], + ids=["bfloat16", "float32"], +) +def test_action_normalization_round_trip_check(lerobot_dataset_metadata, dtype, mean_error): + """The round-trip check the pi05 / pi07 GPU integration tests make, run on CPU in the gate. + + Those tests apply ``assert_action_normalization_round_trips`` to a bfloat16 policy built on + these fixture stats, nightly, on one fresh unseeded draw. Its bound must therefore hold on + every draw. It must still reject an inverse whose mean is slightly off, and on float32, + whose bound is tighter, by less. + """ + stats = lerobot_dataset_metadata.stats + features = {"actions": PolicyFeature(type=FeatureType.ACTION, shape=stats["actions"]["mean"].shape)} + norm_map = {"ACTION": NormalizationMode.MEAN_STD} + + def policy(unnormalize_stats: dict) -> nn.Module: + stub = nn.Module() # the helper reads only these two submodules + stub.normalize_targets = Normalize(features, norm_map, per_dataset_stats=[stats]) + stub.unnormalize_outputs = Unnormalize(features, norm_map, per_dataset_stats=[unnormalize_stats]) + return stub.to(dtype) + + exact, index = policy(stats), torch.zeros(8, dtype=torch.long) + generator = torch.Generator().manual_seed(0) + for scale in (1.0, 4.0): # the GPU tests pass float32 `torch.randn`; 4.0 stretches the tails + for _ in range(20): + actions = torch.randn(8, *features["actions"].shape, generator=generator) * scale + assert_action_normalization_round_trips(exact, actions, index) + + wrong_mean = {**stats, "actions": {**stats["actions"], "mean": stats["actions"]["mean"] + mean_error}} + with pytest.raises(AssertionError): + assert_action_normalization_round_trips(policy(wrong_mean), actions, index) + + # -------------------------------------------------------------------------------------- # State projections read their input dtype off the layer, not off the caller. # -------------------------------------------------------------------------------------- diff --git a/tests/policies/test_pi05.py b/tests/policies/test_pi05.py index 56a243c0..2519c9a6 100644 --- a/tests/policies/test_pi05.py +++ b/tests/policies/test_pi05.py @@ -20,6 +20,7 @@ from opentau.policies.pi05.modeling_pi05 import ( PI05Policy, ) +from tests.utils import assert_action_normalization_round_trips PALIGEMMA_TOKENIZER_EOS_IDS = 1 PALIGEMMA_TOKENIZER_PAD_IDS = 0 @@ -366,15 +367,8 @@ def capture_embed_suffix(*args, **kwargs): # Normalize/Unnormalize now require an explicit `dataset_index`; we're # calling the submodules directly here (bypassing the policy's # `_resolve_dataset_index` helper), so pass a zero-row index manually. - normalize_actions = policy.normalize_targets - unnormalize_actions = policy.unnormalize_outputs - action_output = {"actions": batch["actions"].to("cuda")} dataset_index = torch.zeros(batch_size, dtype=torch.long, device="cuda") - assert torch.allclose( - action_output["actions"], - unnormalize_actions(normalize_actions(action_output, dataset_index), dataset_index)["actions"], - atol=1e-6, - ) + assert_action_normalization_round_trips(policy, batch["actions"].to("cuda"), dataset_index) # Inspect all captured variables expected_vars = [ diff --git a/tests/policies/test_pi07_low_level.py b/tests/policies/test_pi07_low_level.py index 5b1da565..ba4a50f7 100644 --- a/tests/policies/test_pi07_low_level.py +++ b/tests/policies/test_pi07_low_level.py @@ -26,7 +26,7 @@ _global_or_branch_decisions, make_att_2d_masks, ) -from tests.utils import require_vram_gib, to_cuda_bf16 +from tests.utils import assert_action_normalization_round_trips, require_vram_gib, to_cuda_bf16 # VRAM floor for the end-to-end tests below, which build the tiny VLM on CUDA # and run a full forward at 448x448 with T=2. Measured peak on an RTX 5090: @@ -463,15 +463,8 @@ def capture_embed_suffix(*args, **kwargs): # directly bypasses the policy's `_resolve_dataset_index` helper, so # pass an explicit (B,) long index of zeros (this fixture is # single-dataset). - action_output = {"actions": batch["actions"].to("cuda")} dataset_index = torch.zeros(batch_size, dtype=torch.long, device="cuda") - assert torch.allclose( - action_output["actions"], - policy.unnormalize_outputs(policy.normalize_targets(action_output, dataset_index), dataset_index)[ - "actions" - ], - atol=1e-6, - ) + assert_action_normalization_round_trips(policy, batch["actions"].to("cuda"), dataset_index) # Verify all expected captures are present. for var in [ diff --git a/tests/policies/test_pi07_paligemma_low_level.py b/tests/policies/test_pi07_paligemma_low_level.py index f0c70898..03ee2190 100644 --- a/tests/policies/test_pi07_paligemma_low_level.py +++ b/tests/policies/test_pi07_paligemma_low_level.py @@ -27,7 +27,7 @@ PI07PaligemmaLowLevelPolicy, make_att_2d_masks, ) -from tests.utils import require_vram_gib, to_cuda_bf16 +from tests.utils import assert_action_normalization_round_trips, require_vram_gib, to_cuda_bf16 # VRAM floor for the end-to-end test below, which builds the PaliGemma VLM on # CUDA and runs a full forward at 224x224 with T=8. Measured peak in a fresh @@ -466,15 +466,8 @@ def capture_embed_suffix(*args, **kwargs): # directly bypasses the policy's `_resolve_dataset_index` helper, so # pass an explicit (B,) long index of zeros (this fixture is # single-dataset). - action_output = {"actions": batch["actions"].to("cuda")} dataset_index = torch.zeros(batch_size, dtype=torch.long, device="cuda") - assert torch.allclose( - action_output["actions"], - policy.unnormalize_outputs(policy.normalize_targets(action_output, dataset_index), dataset_index)[ - "actions" - ], - atol=1e-6, - ) + assert_action_normalization_round_trips(policy, batch["actions"].to("cuda"), dataset_index) # Verify all expected captures are present. for var in [ diff --git a/tests/utils.py b/tests/utils.py index 322cc6b8..20f696a3 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -151,6 +151,43 @@ def to_cuda_bf16(batch: dict, device: str = "cuda") -> dict: } +def assert_action_normalization_round_trips( + policy: torch.nn.Module, actions: torch.Tensor, dataset_index: torch.Tensor +) -> None: + """Assert that a policy's ``unnormalize_outputs`` inverts its ``normalize_targets`` on ``actions``. + + ``Normalize`` computes in its stats' dtype and casts a floating input to it first, so a + served input gets training's arithmetic. On a policy cast to bfloat16, float32 ``actions`` + therefore come back as bfloat16, exact only to bfloat16's resolution. A plain + ``torch.allclose(actions, round_trip, atol=1e-6)`` both raises on the dtype mismatch and + could never hold. This checks the two properties that do hold: + + * The normalized actions are in the stats' dtype, whatever dtype ``actions`` arrived in. + * The round trip recovers ``actions`` to within the rounding it must incur. MEAN_STD + rounds at most five times (the input cast, ``- mean``, ``/ scale``, ``* scale`` and + ``+ mean``). Each rounding moves the result by at most ``eps / 2`` times + ``|x| + |mean|``. + + A real inversion bug, such as a dropped or wrong mean or a wrong scale, misses by far more + than that. On a float32 policy the bound tightens to float32's ``eps``. + + Args: + policy: Policy whose ``normalize_targets`` / ``unnormalize_outputs`` normalize + ``"actions"`` with ``NormalizationMode.MEAN_STD``. + actions: Unnormalized actions, in any floating dtype. + dataset_index: ``(B,)`` long index of each sample's row in the normalization stats. + """ + stats = policy.normalize_targets.buffer_actions + normalized = policy.normalize_targets({"actions": actions}, dataset_index)["actions"] + assert normalized.dtype == stats["mean"].dtype, "Normalize must compute in its stats' dtype" + restored = policy.unnormalize_outputs({"actions": normalized}, dataset_index)["actions"] + half_ulp = torch.finfo(normalized.dtype).eps / 2 + mean_scale = stats["mean"][dataset_index].abs().max().item() + torch.testing.assert_close( + restored.float(), actions.float(), rtol=5 * half_ulp, atol=5 * half_ulp * mean_scale + ) + + def require_env(func): """ Decorator that skips the test if the required environment package is not installed.