Skip to content

feat(teleop): reset scene objects to their initial poses at end of episode - #190

Open
2047767028-lang wants to merge 1 commit into
AgibotTech:mainfrom
2047767028-lang:feat/reset-scene-after-episode
Open

feat(teleop): reset scene objects to their initial poses at end of episode#190
2047767028-lang wants to merge 1 commit into
AgibotTech:mainfrom
2047767028-lang:feat/reset-scene-after-episode

Conversation

@2047767028-lang

Copy link
Copy Markdown
Contributor

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:

  • Every object has to be put back by hand between episodes. A 150-episode session
    means 150 manual resets.
  • Hand-placed poses are not reproducible. The starting state drifts across the
    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 and
    every articulation after the scene is loaded. task_benchmark.py runs it in
    teleop 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 never
did.

Change

_on_recording() calls _reset_env() right after the /sim/stop_episode signal
has 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_recording already
runs on the render loop (render_step -> _on_recording), so the call is direct
rather 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:

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 scene
reset 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 to
false (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 a
stale path raises inside the render loop. command_controller.py already guards
its own restore path the same way (if not prim.IsValid(): continue), so this
applies 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:

  • Rigid bodies are restored (position + orientation), articulations are restored
    (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.
  • The robot is not reset -- reset_one_frame() skips the "robot" key. The
    arm 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.
  • Velocities are not zeroed, only poses are written. An object that was moving when
    the episode ended keeps its velocity after being teleported back. We tried
    zeroing them by constructing a SingleRigidPrim per object at reset time and it
    invalidated the physics tensor view and shut the simulation down
    (prim '...' was deleted while being used by a tensor view class), so it is left
    out. 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

  • The three touched modules byte-compile.
  • Loading the shipped teleop.yaml through ParameterServer / load_dataclass:
    reset_scene_after_episode reads True; flipping the yaml key reads False;
    removing the key entirely still reads the True default. So the key is really
    wired -- unlike app.auto_start_recording in the same file, which has no
    AppConfig field and is therefore silently ignored (load_dataclass only walks
    dataclass fields; unknown yaml keys disappear without warning -- a warning there
    would have saved us some time).
  • Behaviour was exercised in the simulator on our 3.0.0-based tree: four
    consecutive episodes, each ending with Scene reset to initial state, objects
    back at their initial poses, simulation stable.
  • It has not been run end-to-end against this branch: our environment predates
    the 3.2.0 layout. The change is the same three call sites ported to the new
    paths.

…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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant