[shuheng-liu]

What this does

Fixes three GPU integration tests that have failed on main since #559 (🐛 Bug). This PR changes tests only; the source code is unchanged.

• The failure. Three tests build a policy, cast it to bfloat16, and then check torch.allclose(actions, unnormalize_outputs(normalize_targets(actions)), atol=1e-6) on float32 torch.randn actions:

  • test_pi05.py::TestPI05Integration::test_complete_pi05_pipeline_integration
  • test_pi07_low_level.py::TestPI07LowLevelIntegration::test_complete_pi07_low_level_pipeline
  • test_pi07_paligemma_low_level.py::TestPI07PaligemmaLowLevelIntegration::test_complete_pi07_low_level_pipeline

  fix(serving): keep the state float32 for delta-action targets #559 made Normalize cast a floating input to its stats' dtype first, so served inputs get training's arithmetic. The round trip therefore comes back bfloat16, and allclose raises RuntimeError: Float did not match BFloat16. Comparing in a common dtype is not enough either:

  • atol=1e-6 can never hold at bfloat16 precision;
  • even torch's default bfloat16 tolerances miss near zero, where the fixture's 0.5 mean cancels (up to 0.0022 off).

• The fix. The three identical blocks now call one helper, assert_action_normalization_round_trips, in tests/utils.py next to to_cuda_bf16. It asserts the two properties that hold after fix(serving): keep the state float32 for delta-action targets #559:

  • normalize_targets returns the stats' dtype for a float32 input. This is fix(serving): keep the state float32 for delta-action targets #559's contract, now pinned rather than incidental.
  • The round trip recovers the input within the rounding MEAN_STD must incur. That is at most five roundings (the input cast, - mean, / scale, * scale and + mean), each at most eps / 2 · (|x| + |mean|), with eps read off the stats' dtype. A real inversion bug, such as a dropped or wrong mean or a wrong scale, misses by far more.

• Why the tests change and not the code. fix(serving): keep the state float32 for delta-action targets #559's CHANGELOG entry makes Normalize computing in its stats' dtype intended behavior, and fix(serving): keep the state float32 for delta-action targets #559 left Unnormalize untouched. What went stale is these tests' expectation of a float32-exact round trip on a bfloat16 policy.

• The gate now covers the tolerance. These GPU tests run only nightly. So a new CPU test, test_action_normalization_round_trip_check in tests/policies/test_inference_state_precision.py next to fix(serving): keep the state float32 for delta-action targets #559's own Normalize dtype test, runs the helper on the same fixture stats in bfloat16 and float32 over 40 seeded draws. It also checks that the helper rejects an inverse whose mean is off by 1/16 (bfloat16) or 1e-3 (float32).

• No CHANGELOG entry, since CLAUDE.md exempts test-only PRs.

Sweep. On main, grep -rn "normalize_targets(\|unnormalize_outputs(" tests/ finds exactly these three round trips. test_policies.py::test_backward_compatibility has only a skip-marked case, and no other GPU test round-trips through the normalizers.

How it was tested

• GPU, on a GPU dev box:
  • On main @ a048e55, all three tests fail at their round-trip line with RuntimeError: Float did not match BFloat16.
  • On this branch (c9a42f5), pytest -m "gpu" -n 0 tests/policies/test_pi05.py tests/policies/test_pi07_low_level.py tests/policies/test_pi07_paligemma_low_level.py gives 11 passed and 0 failed. This includes the assertions after the round trip in the pi05 test (pad masks, position ids, attention masks), which had not run on main since fix(serving): keep the state float32 for delta-action targets #559.
• CPU:
  • The new test_action_normalization_round_trip_check passes for bfloat16 and float32.
  • Separately, the helper held on every draw for four stat sets (the fixture's, plus a non-power-of-two std, a negative mean and a tiny std) with inputs scaled up to ×50, in both dtypes: 180 draws of 8 samples per configuration. It caught a 1/16 mean error in every configuration.
• Mutation check (CLAUDE.md rule 8): the CPU test fails when the helper's tolerance is far too tight, far too loose, missing its absolute (mean) term, missing its relative term, or hard-codes bfloat16's eps (so it would not tighten on float32). It also fails when the dtype check is bypassed.
• CPU suite, run the way cpu_test.yml runs it: 3739 passed, 14 skipped. Pre-commit is clean.

How to checkout & try? (for the reviewer)

pytest -sx tests/policies/test_inference_state_precision.py -k round_trip_check

pytest -m "gpu" -n 0 tests/policies/test_pi05.py tests/policies/test_pi07_low_level.py tests/policies/test_pi07_paligemma_low_level.py

Checklist

• I have added Google-style docstrings to important functions and ensured function parameters are typed.
• My PR includes policy-related changes.
  • If the above is checked: I have run the GPU pytests (pytest -m "gpu") and regression tests.

Note: Before submitting this PR, please read the contributor guideline.

[claude]

Reviewed c9a42f5: verified the helper's rounding-error bound (≤ u·(5|x|+3|mean|), inside the asserted 5u·(|x|+max|mean|); the shared std + eps denominator cancels exactly between Normalize and Unnormalize), the MEAN_STD/buffer_actions assumptions against all three GPU test fixtures, and the sweep claim (no other test round-trips the normalizer pair). No issues found.

[claude]

[claude-review] summary for commit c9a42f5

No blocking issues found.