refactor(policies): remove init_strategy

[shuheng-liu]

What this does

Removes the init_strategy config knob entirely from PI0Config, PI05Config, PI05MemConfig, and PI05ContinuousStateConfig, along with the _init_model / _init_weights methods that consumed it.

Why

The knob had three modes — "no_init", "full_he_init", "expert_only_he_init" — but in practice none of them did meaningful work in any shipped config:

• Every shipped config sets pretrained_path. The post_init logic silently overrode init_strategy to "no_init" whenever pretrained_path was set, so "expert_only_he_init" (which appeared in 5 configs including the LIBERO repro) was always dead config.
• "expert_only_he_init" was the only mode that did something the rest of the system couldn't already do — load google/paligemma-3b-pt-224 inside the model constructor for the random-expert + pretrained-backbone use case. We've decided to drop that use case; users who want a pretrained backbone go through pretrained_path exclusively. If someone wants random-expert + pretrained-backbone they can build that checkpoint and upload it themselves.
• That leaves only "no_init" (skip init) and "full_he_init" (He-init everything). When from_pretrained is used, He-init is wasted work because the state dict overwrites it. When fresh-from-scratch is used, no shipped config exercises that path. So the knob is dead either way.

Changes

• Drop init_strategy field, docstring entry, and __post_init__ resolution from all four policy configs. Drop now-unused logging / Literal imports.
• Drop _init_model() call and _init_model / _init_weights method bodies from the four modeling_*.py constructors.
• Drop load_pretrained_paligemma from PaliGemmaWithExpertConfig / PaliGemmaWithExpertModel in pi0 and pi05; PaliGemma is always built from config.
• In policies.py: rename _TRANSIENT_POLICY_FIELDS to _REMOVED_POLICY_FIELDS and call strip_deprecated_fields_from_json(Path(config_file)) before draccus.parse in both PreTrainedConfig.from_pretrained and TrainPipelineConfig.from_pretrained. This is needed because draccus errors on unknown fields, and old saved configs (e.g. TensorAuto/pi05_base/config.json) still contain "init_strategy".
• Strip init_strategy entries from all 9 shipped JSON configs.

Backward-compat note

Old saved configs (HF Hub or local) that still contain "init_strategy" continue to load — the load-time strip drops the key in place before draccus sees it. The next save writes a clean config.

How it was tested

• Each of the four config dataclasses re-loaded with pretrained_path=None, pretrained_path="ckpt", and confirmed no init_strategy attribute exists post-construction.
• Integration test with real draccus: a synthetic JSON containing "init_strategy": "expert_only_he_init" raises DecodingError on direct draccus.parse, but parses cleanly after strip_deprecated_fields_from_json runs first.
• _strip_keys / strip_deprecated_fields_from_json exercised at top-level and nested-under-policy.
• All 9 modified JSON configs round-trip through json.load cleanly.
• All modified Python files compile (python -m py_compile).
• GPU tests not run in this sandbox — flagging for reviewer.

How to checkout & try? (for the reviewer)

git checkout claude/review-weight-init-GQYCz
# Confirm an old-style config (with init_strategy) still loads:
python -c "
import json, tempfile
from pathlib import Path
from opentau.configs.policies import strip_deprecated_fields_from_json
with tempfile.NamedTemporaryFile('w', suffix='.json', delete=False) as f:
    json.dump({'chunk_size': 50, 'init_strategy': 'expert_only_he_init'}, f); p=f.name
strip_deprecated_fields_from_json(Path(p))
print(open(p).read())  # init_strategy stripped
"

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: this PR does touch policy code, but the GPU pytests / regression tests were not run in this sandbox. Leaving the policy box unchecked to reflect that honestly — please flag if a GPU run is required before merge.

https://claude.ai/code/session_01KaAzgu1xy7io6Etf6MXYnJ

[shuheng-liu]

@claude fix the merge conflicts

[claude]

[claude-fix] Merged latest origin/main (8583dc4) into the PR branch in commit b0f5445. Two literal docstring conflicts and one runtime fallout from the merge:

• src/opentau/policies/pi05/paligemma_with_expert.py — kept main's improved attention_implementation docstring (now lists "eager", "sdpa", "fa2", matching the validator). Dropped the load_pretrained_paligemma line main added — that parameter was removed by this PR.
• src/opentau/policies/pi05_mem/configuration_pi05.py — kept main's improved attention_implementation and freeze_vision_encoder docstrings (the latter now correctly describes the SigLIP vision tower, not V-JEPA2). Dropped the init_strategy line main added — that field was removed by this PR.
• tests/policies/test_pi05_mem_gpu.py (new from main) — its PI05MemConfig(init_strategy="no_init", …) would now raise TypeError because the dataclass no longer accepts init_strategy. Removed the kwarg and updated the adjacent comment to reference pretrained_path=None as the equivalent skip-loads signal.

CPU tests: passed locally — pytest -m "not gpu" -n auto --ignore=tests/policies/test_pi07_paligemma_low_level_planner.py tests/policies/ tests/configs/ tests/scripts/test_train.py → 159 passed, 2 skipped. (The pi07 low-level test is the same one CI's cpu_test.yml skips per #209 — pre-existing breakage on main, unrelated to this PR.)

Two follow-up items I deliberately did NOT touch — they may warrant separate PRs:

1. pi06 (added in feat(policies): add authentic pi06 policy with Gemma 3 4B backbone #178) still has its own init_strategy field/logic with all three modes including expert_only_he_init. Cleaning that up to match this PR's premise is out of scope for "fix the merge conflicts", but worth noting since it contradicts the PR title "remove init_strategy". Worse — _REMOVED_POLICY_FIELDS = ("init_strategy",) in policies.py will now silently strip init_strategy from saved pi06 configs at load time, dropping the user's choice on the floor. This is a real regression for pi06 users, and probably wants to be fixed before merge by either (a) dropping init_strategy from pi06 entirely, or (b) making the strip policy-type-aware.
2. pi07_paligemma high/low planners still pass load_pretrained_paligemma=False to PaliGemmaWithExpertConfig. After this PR removes that parameter, the kwarg gets silently absorbed by **kwargs (no crash) but it's dead code. Minor.

Happy to address either or both of those if you'd like — let me know.

[shuheng-liu]

@claude fix the first one in pi06. Leave pi07_paligemma alone. Also make sure the 2 breaking CPU tests skipped by main is also skipped here.

[claude]

[claude-fix] Done — pushed in a595c8e.

1) pi06 — drop init_strategy entirely (matching the rest of this PR). Without this, _REMOVED_POLICY_FIELDS = ("init_strategy",) would silently strip the user's chosen mode from saved pi06 configs at load time, dropping their selection on the floor. Changes:

• src/opentau/policies/pi06/configuration_pi06.py — remove the init_strategy field, its docstring entry, the __post_init__ validation/resolution, and now-unused logging / Literal imports.
• src/opentau/policies/pi06/modeling_pi06.py — remove _init_model() / _init_weights() and the load_pretrained_gemma3 flag derivation in __init__.
• src/opentau/policies/pi06/gemma3_with_expert.py — drop the load_pretrained_gemma3 parameter from Gemma3WithExpertConfig and the corresponding from_pretrained("google/gemma-3-4b-pt") branch in Gemma3WithExpertModel. Gemma 3 is now always built from config (mirrors the pi05 PaliGemma cleanup).
• configs/examples/pi06_training_config.json — strip the "init_strategy": "expert_only_he_init" line so the config round-trips cleanly without going through _REMOVED_POLICY_FIELDS.

Per your instruction, pi07_paligemma is left alone — its load_pretrained_paligemma=False calls are dead but absorbed by **kwargs, harmless.

2) CPU test skips. Confirmed .github/workflows/cpu_test.yml on this branch already matches main verbatim: same three --ignore= (test_planner.py, test_libero_utils.py, test_pi07_paligemma_low_level_planner.py) and the same two --deselect= (test_make_env_async_vector_env, test_make_env_sync_vector_env). My previous local run only applied the pi07 ignore — I've now re-run with the full skip list to match CI.

Tests: passed locally — pytest -m "not gpu" -n auto --ignore=tests/planner/test_planner.py --ignore=tests/utils/test_libero_utils.py --ignore=tests/policies/test_pi07_paligemma_low_level_planner.py --deselect=tests/envs/test_factory.py::TestMakeEnv::test_make_env_async_vector_env --deselect=tests/envs/test_factory.py::TestMakeEnv::test_make_env_sync_vector_env tests/policies/ tests/configs/ tests/scripts/test_train.py → 159 passed, 2 skipped. tests/policies/test_pi06.py alone → 26 passed, 1 deselected. pre-commit run over the changed files passed.

[claude]

Targeted review of #188 — see inline comments and the summary on the PR.

[claude]

[claude-review] summary for commit c6befca

Scope: removal of init_strategy knob across pi0/pi05/pi05_mem/pi06 + load-time strip of removed/deprecated keys. Diff is mostly subtractive and the cleanup is consistent across the four policies.

Findings:

• suggestion — src/opentau/configs/policies.py:376 (and src/opentau/configs/train.py:405) — strip_deprecated_fields_from_json rewrites the source JSON in place. When config_file came from hf_hub_download, that path is a symlink to a content-addressed blob in ~/.cache/huggingface/hub/.../blobs/<sha> — the open(path, "w") follows the symlink and mutates the blob, leaving cache files whose contents no longer match their SHA. For user-supplied local paths, the side-effect rewrite (incl. reformatting) is also surprising. Strip in memory and write to a tempfile.NamedTemporaryFile before handing to draccus.parse. The on-write strip in _save_pretrained is fine.
• suggestion — src/opentau/configs/policies.py:58 — _REMOVED_POLICY_FIELDS is dropped silently, while _DEPRECATED_LATENCY_FIELDS emits a DeprecationWarning via warn_deprecated_latency_fields. A user loading an older config with init_strategy="expert_only_he_init" gets zero signal that their choice was discarded. Extend the warn helper (or add a sibling) so removed fields surface a deprecation warning the same way.
• nit — src/opentau/policies/pi05/configuration_pi05.py:62-63 — stale init_strategy: entry in the class Args: docstring. The field, the Literal import, and the __post_init__ resolution were removed but the docstring entry was missed. (The other three configs were scrubbed cleanly.)
• nit — tests/policies/test_pi05_mem_gpu.py:234 — print("[e2e] device=cuda dtype=bfloat16 init=no_init n_obs_steps=4") still references init=no_init even though the kwarg is gone. Misleading at debug time.

Non-blocking notes:

• src/opentau/policies/pi07_paligemma/{high,low}_level_planner/modeling_pi07_*.py still pass load_pretrained_paligemma=False to the now-renamed config — absorbed by **kwargs, harmless. The PR author chose to defer per the conversation in the comments.
• No new tests covering _strip_keys / _REMOVED_POLICY_FIELDS round-trip in tests/configs/. The PR description describes manual verification but nothing pinned in CI. Worth a small unit test so the strip-on-load behavior doesn't silently regress.
• GPU pytests not run in the sandbox per the PR description; reasonable, but the test_pi05_mem_gpu.py change does need a GPU runner to actually exercise it.

[shuheng-liu]

@claude fix the suggestions and nits if they make sense

[claude]

[claude-fix] Done — pushed in 4505fac. Walked through the two suggestions and two nits from the review summary in comment 4348009825; replied inline on the two suggestions and addressed the two nits + the test gap here:

• suggestion (in-place rewrite of HF cache symlinks) — addressed in discussion r3164662608. New load_stripped_config_to_tempfile() strips in memory and writes to a tempfile.mkstemp temp file at both load sites; the source path is never mutated.
• suggestion (silent strip of removed fields) — addressed in discussion r3164663068. New warn_removed_policy_fields() emits a DeprecationWarning before the strip at both load sites.
• nit (pi05/configuration_pi05.py:62-63 stale init_strategy: docstring entry) — dropped the two-line entry. The other three configs were already clean.
• nit (tests/policies/test_pi05_mem_gpu.py:234 misleading init=no_init in the e2e debug print) — changed to pretrained_path=None, matching the actual signal that picks the skip-loads path.
• non-blocking note (no CI test for _strip_keys / _REMOVED_POLICY_FIELDS round-trip) — added 7 unit tests in tests/configs/test_utils_policies.py:
  • test_strip_keys_top_level_and_nested / test_strip_keys_no_change — base helper
  • test_load_stripped_config_to_tempfile_does_not_mutate_source — source byte-for-byte unchanged
  • test_load_stripped_config_does_not_follow_symlink — explicitly pins HF-cache safety (creates a symlink to a "blob", strips through the symlink, asserts the linked-to file is unchanged)
  • test_strip_deprecated_fields_from_json_in_place_for_owned_files — the on-write path still works
  • test_warn_removed_policy_fields_emits_deprecation_warning / test_warn_removed_policy_fields_silent_when_clean — new warn helper
• non-blocking note (pi07_paligemma still passes load_pretrained_paligemma=False) — left alone per your earlier instruction in comment 4347923227.

Tests: pytest -m "not gpu" -n auto --ignore=tests/planner/test_planner.py --ignore=tests/utils/test_libero_utils.py --ignore=tests/policies/test_pi07_paligemma_low_level_planner.py --deselect=tests/envs/test_factory.py::TestMakeEnv::test_make_env_async_vector_env --deselect=tests/envs/test_factory.py::TestMakeEnv::test_make_env_sync_vector_env tests/ → 720 passed, 16 skipped in 127s. (3 errors in tests/utils/test_hub.py::test_save_pretrained require HUGGINGFACE_HUB_TOKEN which CI provides but the sandbox does not — unrelated to this PR.) pre-commit run over the changed files passed.