[shuheng-liu]

What this does

Fixes served joint targets for delta-action checkpoints (🐛 Bug).

• Delta-action targets carried bf16 rounding. Every inference entry point (the gRPC and RoboCasa servers, preprocess_observation for eval.py and in-training sim eval, and the dummy observations of inference.py / benchmark_inference.py / high_level_planner_inference.py) built state in the bfloat16 serving dtype. A delta-action checkpoint (delta_action_state_map set) adds that state back onto its predicted deltas, but the dataset formed the training deltas against the float32 state. So every absolute target carried the state's bf16 rounding: up to 2**-7 ≈ 0.0078 rad (0.45°) on a joint in [2, 4) rad, re-applied on every chunk. state and the gRPC RTC prefix_action are now built at INFERENCE_STATE_DTYPE (float32).
• The model's input is unchanged. Training batches carry a bf16 state: BaseDataset._to_standard_data_format casts every float to bf16 after forming the deltas, and Normalize then works in bf16 arithmetic. A float32 state would have normalized in float32 and moved a discretized state token by a bin for 16–60% of states, depending on the stats. So Normalize now casts a floating input to its stats' dtype. That is a no-op in training, and a served policy's normalized input stays bit-identical.
• State projections (pi0, pi05, pi05_mem, and the pi07 / pi07_paligemma low-level policies) now cast to the projection's own weight dtype (cast_to_weight_dtype) instead of a hard-coded bf16. pi0 had no cast at all, so a float32 state crashed it.
• add_chunk_start_state / subtract_chunk_start_state now compute in the wider of their two input dtypes.
• scripts/actions_mse_loss.py scored a delta checkpoint's absolute predictions against delta targets; it now compares in delta space. This also fixes three older bugs:
  • its chunk indexing, which never worked for chunk_size > 1;
  • a KeyError on datasets whose action column is named action;
  • r2_score called with its arguments swapped.

Served actions change only for delta-action checkpoints, plus xr1 / flux3_action, which leave state unnormalized and now preprocess it at float32, as their reference pipelines do. Both fixes have CHANGELOG entries under [Unreleased].

How it was tested

• tests/policies/test_inference_state_precision.py drives the real RobotPolicyServicer.GetActionChunk and the RoboCasa infer / infer_batch against a real PI05Policy, with only the Hub-loaded pieces stubbed. It asserts:

  • served absolute targets equal float32 state + delta, bit for bit;
  • an RTC prefix reaches the model as training's delta prefix;
  • discrete-state tokens and the continuous state input match the training-path values.

  It also covers Normalize, cast_to_weight_dtype and pi0, plus an AST sweep that fails if any script builds state at another dtype.

• tests/scripts/test_actions_mse_loss.py runs a perfect-policy stub on a tiny delta dataset. There are also additions to tests/datasets/test_action_indexing.py and tests/envs/test_utils.py.

• Mutation check: reverting each piece of the fix fails exactly the tests that pin it.

• CPU suite, run as cpu_test.yml runs it: 3699 passed, 14 skipped. Pre-commit is clean.

• Determinism (CLAUDE.md rule 3): configs/dev/ci_config.json, 3 steps on CPU, same seed. Per-micro-batch losses (exact float hex) were bit-identical across two runs of this branch and one run of the base commit. With --policy.state_type=continuous, the first forward loss was bit-identical to the base commit.

• GPU pytests (pytest -m "gpu") and regression tests have not been run. They are deferred to the nightly gpu_test.yml / regression_test.yml, and the GPU checkbox below stays unticked until they pass.

How to checkout & try? (for the reviewer)

pytest -sx tests/policies/test_inference_state_precision.py tests/scripts/test_actions_mse_loss.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]

[claude-review] summary for commit 3807c6f

No blocking issues found.