[shuheng-liu]

What this does

Fixes a log flood in every discrete-state policy (🐛 Bug). Only the logging changes: the state tokens stay bit-identical.

• The warning fired on every call. prepare_discrete_state clamps the normalized state into [-1, 1], then bins it into 256 prompt tokens. It logged State values are not normalized between -1 and 1 whenever any value fell outside that range. Under QUANTILE / MIN_MAX normalization a real robot's state routinely overshoots a little, because q01/q99 and the dataset min/max are not hard bounds. So the gRPC server logged the warning on nearly every request. In a real deployment this flooded the log, and on a nearly full disk the synchronous writes stalled individual requests for seconds.
• One shared helper now rate-limits it. The helper is warn_once_if_state_out_of_range(state, policy_name) in src/opentau/policies/utils.py.
  • It logs the WARNING, with that batch's min and max, once per process for each policy (keyed by the policy's name).
  • After that it only counts. An INFO line after the 10th, 100th, 1000th, ... out-of-range call gives the count and the extremes so far, so a persistent or worsening overshoot stays visible at logarithmic cost.
  • The tally sits behind a lock for multi-threaded servers, but the logging happens outside the lock.
  • The helper is @torch.compiler.disabled. If traced, its data-dependent branch and .item() calls would split a compiled sample_actions into several small graphs (measured: 5 graphs and 3 breaks, against 1 graph and 1 break with the decorator).
• All five copies route through it: pi05 (which pi05_ttt inherits), pi06, the pi07 high-level planner, the pi07_paligemma high-level planner, and value. The clamp and binning lines are untouched. In value's copy, state.to(...).numpy() is split into two statements so the helper receives the tensor.
• Docstrings fixed (CLAUDE.md rule 8). Four copies documented Raises: ValueError: If the state values are not normalized between -1 and 1., but every copy only warned and clamped. All five now say that out-of-range values are clamped into [-1, 1] (bin 0 or 255) and logged once per process.
• CHANGELOG: one ### Fixed entry appended to [Unreleased]. The change is logging-only, but it changes what a serving process logs and removes the stalls. CLAUDE.md exempts only docs-only, test-only and chore(claude) PRs from an entry.

Sweep and inclusion criterion. The copies were enumerated mechanically before scoping. grep -rn "State values are not normalized" src/ and grep -rn "def prepare_discrete_state" src/ both return exactly these five.

• Included: every prepare_discrete_state that logs about the out-of-range state it clamps. All five warned and clamped, and none raised, so each one keeps that behavior.
• Excluded: scripts/onnx_inference.py::_prepare_discrete_state. It clamps without logging, so there is nothing to rate-limit.
• A subclass that inherits a copy (pi05_ttt) is covered by its base. A new test enforces this list (see below).

How it was tested

• New tests/policies/test_discrete_state_out_of_range.py (13 CPU tests). The tests call each copy's real, unbound prepare_discrete_state with a stub self. The method reads only name, plus config.state_type in pi05, so no model is built. The tests check that:
  • for each of the five copies, 5 calls with an out-of-range state log exactly one WARNING (checked with caplog), and every call, warned or silent, returns the hand-computed clamped bins;
  • a value of exactly ±1 does not warn, and in-range calls don't use up the one warning;
  • in one process, each policy warns once and none silences another;
  • INFO summaries fire at the 10th, 100th and 1000th out-of-range call with the right counts and running extremes, and in-range calls are not counted;
  • in an AST scan of the package, every def prepare_discrete_state is in the test's list and calls the helper, and none calls logging.* directly. A sixth copy therefore can't silently bring back the per-call warning.
• Mutation check (rule 8): each of the following mutations fails at least one test:
  • warning on every call;
  • one global flag instead of per-policy keys;
  • counting in-range calls;
  • one copy going back to an inline warning;
  • a dropped clamp;
  • an off-by-one bin;
  • treating exactly ±1 as out of range.
• Bit-identical output: main's and this branch's prepare_discrete_state ran for all five copies on ~1.9M values: uniform and normal float32, bfloat16, every bin edge ±1 ulp, ±1e30 and ±inf. The outputs are identical.
• Determinism (rule 3): ran configs/dev/ci_config.json (pi05, discrete state) for 3 steps on CPU with the same seed. The 6 per-micro-batch losses (exact float hex of loss, MSE and CE) were bit-identical across two runs of this branch and one run of the base commit (main @ a048e55). This config min-max-normalizes the state with the dataset's own min/max, so the state never leaves [-1, 1] and the warning path stays idle. A second scenario switched the state to MEAN_STD in a scratch copy of the config, so the state does leave [-1, 1]. The losses were again bit-identical across two branch runs and the base run. The base logged the warning on 5 of the 6 micro-batches; this branch logged it once.
• CPU suite, run the way cpu_test.yml runs it: 3749 passed, 14 skipped, 1 failed. The failure was tests/utils/test_utils_benchmark.py::test_context_manager_usage, a timing assertion where sleep(0.05) measured 60 ms under full-core load. It passes when run on its own. Pre-commit is clean.
• GPU pytests: targeted run of pytest -m "gpu" -n 0 over the five affected policies' test files on a GPU dev box. On commit f96689d: 5 passed, 1 skipped, 1 failed.
  • The failure is pre-existing, and test(policies): compare the normalize round trip at bfloat16 precision #561 fixes it. test_pi05.py::TestPI05Integration::test_complete_pi05_pipeline_integration fails with RuntimeError: Float did not match BFloat16 in its Normalize→Unnormalize round-trip assertion. It fails the same way on main @ a048e55 without this change: fix(serving): keep the state float32 for delta-action targets #559's stats-dtype Normalize met a stale float32 round-trip assertion. Its policy.forward, which runs through the new helper and logs the warning once, passes before that assertion. With test(policies): compare the normalize round trip at bfloat16 precision #561's test fix applied on top of this branch, the same targeted run gives 6 passed, 1 skipped and 0 failed.
  • The skip is the pi07_paligemma planner's GPU test, which is unconditionally @pytest.mark.skipped. That copy is covered by the CPU tests above.
  • Not run: the nightly regression suite (regression_test.yml).

How to checkout & try? (for the reviewer)

pytest -sx tests/policies/test_discrete_state_out_of_range.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 commit f96689d. No blocking issues found — details in the [claude-review] summary comment.

Verified locally: all 13 new tests in tests/policies/test_discrete_state_out_of_range.py pass, and the full tests/policies CPU subset passes (1737 passed, 2 skipped). The five-copy sweep is complete (grep -rn "def prepare_discrete_state" src/ returns exactly the five routed copies plus the non-logging onnx_inference.py helper), all five policy name keys are distinct, and the clamp/binning lines are untouched.

[claude]

[claude-review] summary for commit f96689d

No blocking issues found.

Verification performed on this run:

• Ran tests/policies/test_discrete_state_out_of_range.py: 13 passed.
• Ran pytest -q tests/policies -m "not gpu and not network" -n auto: 1737 passed, 2 skipped.
• Confirmed the sweep is complete: grep -rn "def prepare_discrete_state" src/ returns exactly the five routed copies plus the intentionally excluded non-logging scripts/onnx_inference.py::_prepare_discrete_state; pi05_ttt has no override and inherits pi05's copy.
• Confirmed all five policy name attributes are distinct (pi05, pi06, pi07_high_level, pi07_high_level_planner, value), so the per-policy rate-limit keying works.
• Checked the helper: calls == 1 is handled before the power-of-10 elif, so the first call logs exactly one WARNING; the tally mutation is inside the lock and logging outside it; value's .to(...).numpy() split is aliasing-equivalent, and the clamp/binning lines are untouched in all five copies.
• CHANGELOG entry is appended at the end of [Unreleased] with no double blank lines; no other code or test references the old warning text; logging/np imports remain used in every modified file.