refactor!: state.info ownership — episode state on env buffers, typed reward_terms, reset -> None - #59
Conversation
ddb22b8 to
5868ad5
Compare
5868ad5 to
ab12449
Compare
There was a problem hiding this comment.
🔵 Needs a closer look
This is a large breaking refactor spanning ~50 files, and it already introduced a confirmed reward regression (dead grasp_success_bonus via buffer aliasing in insert_peg), so it needs human review before approval.
Pull request overview
This PR completes the removal of the state.info channel across the framework. Episode-scoped task state that previously lived in state.info[...] is moved onto env-instance buffers (self._xxx, allocated full-batch in __init__ and rewritten in place at reset). The env state dataclasses (ArrayEnvState / TorchEnvState) drop info in favor of a typed reward_terms: dict[str, ndarray] field alongside metrics, reset() now returns None, and the RL boundary is composed by env_infos() as {"Reward": state.reward_terms, "metrics": ...}. The RSLRL wrappers derive time_outs from truncated & ~terminated rather than reading it from info.
Changes:
- Removed
state.info; added typedreward_termsto the core state dataclasses and madereset()returnNone, deleting the reset-info merge machinery. - Migrated per-env episode state (commands, phases, grasp/gripper trackers, action-delay buffers, etc.) from
infodict entries to full-batchself._xxxbuffers across ~15 environments. - Rewired RL wrappers/
env_infos()and derivedtime_outsfrom lifecycle flags; updated tests and bilingual docs accordingly.
File summaries
| File | Description |
|---|---|
| motrix_env_core/.../array/env.py | Adds reward_terms field, reset() -> None, removes info/merge helpers |
| motrix_env_motrixsim/.../torch_env.py | Same state-contract change for the Torch frontend |
| motrix_rl/src/motrix_rl/utils.py | env_infos() now returns {"Reward": reward_terms, "metrics": ...} |
| motrix_rl/.../skrl/{torch,jax}/wrap_{np,torch}.py | reset() returns env_infos(state) instead of state.info |
| motrix_rl/.../rslrl/torch/wrap_{np,torch}.py | Derive time_outs from truncated & ~terminated |
| motrix_rl/.../fastsac/wrap_{np,torch}.py | last_info sourced from env_infos(state) |
| motrix_rl/.../deploy/source_rollout.py | Pin commands via env._commands buffer |
| motrix_envs/.../rm65_insert_peg/insert_peg_np.py | Grasp/tracking migrated to buffers — aliasing regression flagged |
| motrix_envs/.../ (stewart, quadruped walk, shadow_hand, franka×2, rm65×2, go1, bounce_ball, etc.) | Episode state migrated from info to self._xxx buffers |
| motrix_envs/tests/, motrix_rl/tests/ | Tests updated to read buffers / reward_terms instead of info |
Review details
- Files reviewed: 67/67 changed files
- Comments generated: 1
- Review effort level: Balanced
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
ab12449 to
c1e6082
Compare
There was a problem hiding this comment.
🔵 Needs a closer look
It is a wide-reaching breaking refactor touching the core state contract and every environment's in-place buffer semantics, where subtle aliasing regressions cannot all be exhaustively verified, so it warrants final human review.
Review details
Suppressed comments (1)
motrix_envs/src/motrix_envs/motion/sampler.py:91
- Same garbled replacement as the class docstring:
motion-step state[state.terminated]is not a valid or meaningful reference after thestate.inforemoval. Point it at the actual motion-step buffer indexed by the failed rows (cf.wbt/mdp/command.py:164).
- Files reviewed: 67/67 changed files
- Comments generated: 1
- Review effort level: Balanced
c1e6082 to
767356b
Compare
There was a problem hiding this comment.
🔵 Needs a closer look
A sweeping breaking contract change across the core and ~15 environments with subtle aliasing/partial-reset semantics, including a confirmed latent name-collision bug in rm65_open_cabinet, warrants human review.
Review details
- Files reviewed: 67/67 changed files
- Comments generated: 1
- Review effort level: Balanced
…outs derived in RSLRL, info replaced by typed reward_terms Complete cleanup of the state.info channel in three moves: 1. Framework info ownership: ArrayEnv/TorchEnv stop maintaining info["time_outs"]; the RSLRL wrappers derive it from truncated & ~terminated for value bootstrapping (SKRL already used state.truncated, fastsac its own terminated/truncated returns). stewart stops overwriting the framework flag with a local float copy. 2. Episode-scoped task state migrates from info dicts to env-instance full-batch buffers across all environments (anymal_c, stewart, shadow_hand, franka x2, rm65 x2, quadruped walk, go1, bounce_ball, manipulator, basic quadruped, finger, pendulum): __init__ allocates the buffers, reset writes done rows in place, consumers read them directly with no info.get()/membership fallbacks. Dead writes removed (target_info, phase2_mask, handle_pose_override path, max_consecutive_bounces, action_delta, ...). Fixes partial-reset shape bugs in stewart's disturbance writes and walk_np's _randomize_params, and hopper's hop branch never writing its reward terms. 3. state.info is removed: ArrayEnvState/TorchEnvState gain a typed reward_terms field for the per-term reward breakdown (freshly written every transition) alongside metrics; reset() -> None so the reset-info merge machinery (_merge_reset_info/_replace_info_values) is deleted; the manager frontend writes reward_terms from the kernel layout; motrix_rl's env_infos() composes the RL-boundary dict from the typed fields; deploy source_rollout pins commands via the env buffer (silently broken since the walk_np migration). Reset-time Reward zeroing dropped (quadruped's three static init tables included). Tests and bilingual docs updated. Episode state lives on env buffers, lifecycle flags and reward terms on the typed state object, reduced diagnostics in metrics; info no longer exists.
767356b to
fbfc4bd
Compare
There was a problem hiding this comment.
🔵 Needs a closer look
It is a wide-reaching breaking refactor of the env/RL state contract across ~40 files with subtle in-place buffer semantics that already produced (now-fixed) bugs, warranting final human verification.
Review details
- Files reviewed: 67/67 changed files
- Comments generated: 0 new
- Review effort level: Balanced
Summary
Complete cleanup of the
state.infochannel, in three moves:Framework info ownership —
ArrayEnv/TorchEnvstop maintaininginfo["time_outs"]; the RSLRL wrappers derive it fromtruncated & ~terminatedfor value bootstrapping (SKRL already usedstate.truncated, fastsac its ownterminated/truncatedreturns). stewart stops overwriting the framework flag with a local float copy.Episode-scoped task state → env-instance buffers — across all environments (anymal_c, stewart, shadow_hand, franka ×2, rm65 ×2, quadruped walk, go1, bounce_ball, manipulator, basic quadruped, finger, pendulum):
__init__allocates full-batch buffers, reset writes done rows in place, consumers read directly with noinfo.get()/membership fallbacks.state.inforemoved —ArrayEnvState/TorchEnvStategain a typedreward_terms: dict[str, ndarray]field (per-term reward breakdown, rewritten every transition) alongsidemetrics;reset() -> None, so the reset-info merge machinery (_merge_reset_info/_replace_info_values) is deleted; the numba manager frontend writesreward_termsfrom the kernel layout;motrix_rl.env_infos()composes the RL-boundary dict from the typed fields (trainers unchanged).Final contract
self._xxx)state.terminated/state.truncated/state.episode_stepsstate.reward_termsstate.metricsenv_infos()→{"Reward": ..., "metrics": ...}time_outsBugs fixed along the way
_apply_disturbance_to_stage: full-width compute vs subset write on partial auto-reset (shape mismatch)_randomize_params: latent partial-reset shape bug via first-insert mergedeploy/source_rolloutcommand pinning silently broken since the walk_np buffer migrationtarget_info,phase2_mask, always-identity_resolve_handle_poseoverride path,max_consecutive_bounces,action_delta, reset-time Reward zeroing (incl. quadruped's three staticinit_reward_infotables, ~70 lines)Breaking changes
state.infono longer exists; external consumers must usestate.reward_terms/state.metricsreset(env_ids)returnsNone;TorchEnv.reset(data)returns observations onlyinfo["time_outs"]must derivetruncated & ~terminatedTest plan
rsl_rl/tensordict, verified on base)state.reward_terms/state.metrics)