feat(teleop): reset scene objects to their initial poses at end of episode - #190
Open
2047767028-lang wants to merge 1 commit into
Open
feat(teleop): reset scene objects to their initial poses at end of episode#1902047767028-lang wants to merge 1 commit into
2047767028-lang wants to merge 1 commit into
Conversation
…isode Multi-episode collection (AgibotTech#182) lets the operator record episode after episode without restarting the simulator, but the scene is never restored between them: whatever the last episode moved, picked up or knocked over stays where it ended. The operator has to put every object back by hand between episodes -- 150 episodes means 150 manual resets -- and hand-placed poses are not reproducible, which shows up later as an inconsistent starting distribution in the training data. The capability already exists, it was simply never called on the collection path: - collect_init_physics() snapshots the initial world pose of every rigid body and every articulation after the scene is loaded, and task_benchmark.py runs it in teleop mode too; - reset_one_frame() writes those poses back; - _reset_env() ties the two together. Only the evaluation path (pi_env.py) called it. Teleop collection never did. This change calls _reset_env() from _on_recording() right after the /sim/stop_episode signal has been handled, i.e. once the bag is closed and the one-shot latches are re-armed, so the next episode starts from the recorded initial state. _on_recording already runs on the render loop, so the call is direct rather than queued through run_on_render_loop. Scope: this only fires on /sim/stop_episode, which is the multi-episode signal added in AgibotTech#182. A workflow that records a single episode per simulator launch never reaches this code and is unaffected. Also guards the rigid-body branch of reset_one_frame() with a prim validity check. An object captured at init time can be gone by the time the episode ends (the API exposes delete_prim), and touching a stale path raises inside the render loop. command_controller.py already guards its own restore path the same way. reset_scene_after_episode (AppConfig + teleop.yaml) turns it off. It defaults to true, so multi-episode collection resets by default; note this is a behaviour change for existing multi-episode users, one line of config to opt out. Verified: the three touched modules byte-compile; loading the shipped teleop.yaml through ParameterServer/load_dataclass yields reset_scene_after_episode=True, flipping the yaml key yields False, and removing the key entirely still yields the True default (i.e. the key is really read -- unlike app.auto_start_recording, which has no AppConfig field and is silently ignored). Behaviour was exercised in the simulator on our 3.0.0-based tree, where the same change reset the scene on four consecutive episodes; it has not been run end-to-end against this branch, since our environment predates the 3.2.0 layout. Co-Authored-By: Claude Opus 5 <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.
Problem
With multi-episode collection (#182) the operator records episode after episode
without restarting the simulator, but the scene is never restored between them:
whatever the last episode moved, picked up or knocked over stays where it ended.
Two consequences:
means 150 manual resets.
session, which shows up later as an inconsistent initial-state distribution in
the training data -- the operator cannot tell how far it drifted, because there
is no record of where things were supposed to be.
The capability already exists, it was just never called here
collect_init_physics()snapshots the initial world pose of every rigid body andevery articulation after the scene is loaded.
task_benchmark.pyruns it inteleop mode too, so the snapshot is available during collection.
reset_one_frame()writes those poses back._reset_env()ties the two together.Only the evaluation path (
pi_env.py) called it. The teleop collection path neverdid.
Change
_on_recording()calls_reset_env()right after the/sim/stop_episodesignalhas been handled -- once the bag is closed and the one-shot latches are re-armed --
so the next episode starts from the recorded initial state.
_on_recordingalreadyruns on the render loop (
render_step -> _on_recording), so the call is directrather than queued through
run_on_render_loop.reset_scene_after_episode(AppConfig+teleop.yaml) turns it off.Relationship to multi-episode recording (#182)
This is meant to be used together with multi-episode recording, and only works in
that flow:
/sim/stop_episode, the signal introduced in feat(teleop): multi-episode recording without restarting the simulator #182. A workflowthat records one episode per simulator launch never reaches this code, and gets
the process-level reset from the restart anyway.
the next episode, but the scene is not. That is the gap this closes.
So the pair is "the simulator is reusable across episodes" (#182) + "the scene is
too" (this PR). Worth stating in the docs if the multi-episode flow is written up.
Behaviour change
The default is
true, so existing multi-episode users will start seeing the scenereset between episodes where it previously stayed put. That is the behaviour we
believe is wanted for collection, but it is a change -- one line of config
(
app.reset_scene_after_episode: false) opts out. Happy to flip the default tofalse(opt-in) if you would rather not change existing behaviour.Also included: prim validity guard in
reset_one_frame()An object captured at init time can be gone by the time the episode ends (the API
exposes
delete_prim, and scenes can be rebuilt between episodes). Touching astale path raises inside the render loop.
command_controller.pyalready guardsits own restore path the same way (
if not prim.IsValid(): continue), so thisapplies the same pattern. Four lines, no behaviour change when all prims are valid.
What is and is not reset
Worth knowing before relying on it:
(joint state + base pose). Which objects those are is discovered by traversing
the stage in
collect_physics(), so this is not tied to any particular scene.reset_one_frame()skips the"robot"key. Thearm stays where the operator left it, so they still need to return it to the
start pose before recording the next episode. Left as is, since changing it
would affect the evaluation path that already uses this function.
the episode ended keeps its velocity after being teleported back. We tried
zeroing them by constructing a
SingleRigidPrimper object at reset time and itinvalidated the physics tensor view and shut the simulation down
(
prim '...' was deleted while being used by a tensor view class), so it is leftout. If you want it, the safe form is the one in
command_controller.py:build the rigid-body objects once at scene load and reuse them.
Verification
teleop.yamlthroughParameterServer/load_dataclass:reset_scene_after_episodereadsTrue; flipping the yaml key readsFalse;removing the key entirely still reads the
Truedefault. So the key is reallywired -- unlike
app.auto_start_recordingin the same file, which has noAppConfigfield and is therefore silently ignored (load_dataclassonly walksdataclass fields; unknown yaml keys disappear without warning -- a warning there
would have saved us some time).
consecutive episodes, each ending with
Scene reset to initial state, objectsback at their initial poses, simulation stable.
the 3.2.0 layout. The change is the same three call sites ported to the new
paths.