test(policies): compare the normalize round trip at bfloat16 precision

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.
 # --------------------------------------------------------------------------------------

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 = [

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 [

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 [

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.