Add environment task record-and-replay (async parallel envs + auto-save)#425
Open
yuecideng wants to merge 20 commits into
Open
Add environment task record-and-replay (async parallel envs + auto-save)#425yuecideng wants to merge 20 commits into
yuecideng wants to merge 20 commits into
Conversation
Design a record-and-replay feature for EmbodiChain environments with two modes: pure kinematic (disable physics, write recorded object states each step) and dynamic (feed recorded robot trajectory, re-simulate physics). Reuses the existing expert-mode rollout_buffer (extended with a `states` field), the _hook_after_sim_step write hook, and the gym.Wrapper pattern. Co-Authored-By: Claude <noreply@anthropic.com>
Six TDD tasks: (1) trajectory states buffer builder + loader, (2) record kinematic states into rollout_buffer, (3) save_trajectory persistence, (4) base_env auto-reset guard, (5) ReplayWrapper pure-kinematic mode, (6) ReplayWrapper dynamic mode + edge cases. Co-Authored-By: Claude <noreply@anthropic.com>
Co-Authored-By: Claude <noreply@anthropic.com>
Co-Authored-By: Claude <noreply@anthropic.com>
Co-Authored-By: Claude <noreply@anthropic.com>
…ings, sanity check, test)
yuecideng
added
enhancement
yuecideng
marked this pull request as ready for review
July 25, 2026 06:57
Dedicated per-env trajectory buffer (states + pre-process action) with
per-env step counters to support async parallel envs; default auto-save
to ~/.cache/embodichain_data/trajectories/{run_id}/; multi-env replay
with per-env lengths. Pre-process action recording enables ActionManager
in dynamic replay. Supersedes the recording/replay portions of the
original spec.
Co-Authored-By: Claude <noreply@anthropic.com>
Four TDD tasks: (1) ReplayWrapper per-env + pre-process action
(format-compatible); (2) dedicated _traj_buffer with per-env recording +
save; (3) default auto-save to ~/.cache/embodichain_data/trajectories/
{run_id}/; (4) async + ActionManager dynamic-replay tests. Sequenced so
replay is bi-directionally format-compatible, keeping each task green.
Co-Authored-By: Claude <noreply@anthropic.com>
…ocess action Co-Authored-By: Claude <noreply@anthropic.com>
…_action Co-Authored-By: Claude <noreply@anthropic.com>
- Add test_async_envs_do_not_corrupt_recording proving per-env _traj_steps isolates envs when one resets early. - Add ReplayDeltaEnv (delta-qpos ActionManager) and test_dynamic_replay_with_action_manager proving raw pre-process actions are recorded and dynamic replay tracks them. Co-Authored-By: Claude <noreply@anthropic.com>
…ype hints Co-Authored-By: Claude <noreply@anthropic.com>
yuecideng
changed the title
run_env now supports three modes: data generation (default), preview (--preview), and replay (--replay). Replay loads a .pt trajectory and drives it via ReplayWrapper in kinematic (exact), dynamic (physics), or control (interactive kinematic scrubber via terminal input) mode. - ReplayWrapper.go_to_step(step): O(1) scrub to an arbitrary recorded step (clamped to min(lengths)-1); powers control mode. - run_env: replay()/replay_auto()/replay_control() + --replay, --replay_trajectory, --replay_mode args; --replay and --preview are mutually exclusive. - Tests: go_to_step scrubs to the recorded state (forward/clamp/back); run_env.replay() runs kinematic + dynamic without error. Co-Authored-By: Claude <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
This PR adds a record-and-replay capability to the environment framework. An episode can be saved as a trajectory file (
.pt) and replayed later in one of two modes:env.stepso physics re-simulates; produces the fullobs/reward/terminated/truncated/info.Key capabilities
self._traj_buffer(states + actions) with a per-env step counter (_traj_steps) means envs that terminate/reset at different times no longer corrupt each other's recording. Only the reset envs' counters clear.actionis the raw action passed toenv.step(before_preprocess_action), so dynamic replay feeds it back throughenv.stepand theActionManagerre-applies the same transform - dynamic replay is faithful even with an ActionManager (delta/eef_pose action terms).~/.cache/embodichain_data/trajectories/{run_id}/at episode end and onclose()(best-effort: IO errors warn + skip, never crash the episode). Explicitsave_trajectory(path, env_ids=None)also available.ReplayWrapperreplays per-env lengths (meta["lengths"]); finished envs hold their last state/action. Backward-compatible with older files (nolengths-> uniform length).It reuses the existing
gym.Wrapperpattern, the_hook_after_sim_stephook, the auto-reset guard inbase_env.step, and the existing per-object state read/write APIs. The sharedrollout_buffer(obs/actions/rewards),current_rollout_step, LeRobot recorder, and RL mode are untouched - the trajectory buffer is fully decoupled. All new behavior defaults off (record_trajectory=False), so existing envs are unaffected.Design spec & implementation plan:
docs/superpowers/specs/2026-07-24-env-replay-async-design.mdanddocs/superpowers/plans/2026-07-24-env-replay-async.md(developed via brainstorming -> writing-plans -> subagent-driven-development, with per-task reviews and a final whole-branch review).Components
build_trajectory_buffer/load_trajectory- pure helpers inembodichain/lab/gym/utils/gym_utils.py.EmbodiedEnvCfg.record_trajectory/trajectory_uids/trajectory_save_dir/trajectory_auto_savecfg fields.EmbodiedEnv._traj_buffer+_traj_steps(per-env);_preprocess_actionstashes the raw action;_write_trajectory_steprecords states + action per-env.EmbodiedEnv.save_trajectory(path, env_ids=None)+_save_trajectory_for_env(auto-save).BaseEnv.stepauto-reset guard (_replay_no_auto_reset) so dynamic replay isn't disrupted mid-trajectory.ReplayWrapper(gym.Wrapper)- kinematic + dynamic modes, per-env lengths, num_envs broadcast, physics/flag restoration onclose(), active_joint_ids/robot_dof sanity check.Scope (deferred)
"rl"-mode rollout buffer /set_rollout_bufferinjection path.qvelrecording.Dependencies: none new (uses existing
tensordict,torch,gymnasium).Type of change
Checklist
black .command to format the code base.docs/superpowers/; Sphinx user docs are follow-up)Testing
tests/gym/utils/test_gym_utils.py- unit tests for the builder/loader (mock stubs, no sim).tests/gym/envs/test_replay.py- real DexSim (CPU) integration tests: recording populates_traj_buffer; save/load round-trip with per-envlengths; auto-reset guard; kinematic round-trip (state readback == recorded at atol 1e-4); dynamic replay (qpos tracks recorded); per-env truncation timing; legacy-file backward compat; num_envs broadcast; physics/flag restoration; async parallel envs don't corrupt recording; ActionManager dynamic replay (records raw delta, re-applies via ActionManager); auto-save at episode end + close.tests/gym/envs/test_base_env.pyCPU regression still passes (auto-reset guard preserves normal behavior).33 tests pass (13 replay + 19 gym_utils + 1 base_env CPU).