Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Adds a script for recording training data from a Gym simulation and replaying in sim #437

Open
wants to merge 28 commits into
base: main
Choose a base branch
from

Conversation

samzapo
Copy link

@samzapo samzapo commented Sep 12, 2024

What this does

Adds a script (lerobot/scripts/record_and_replay_teleop_dataset.py) that:

  • dynamically imports a named module that implements a gym environment, gym-drake-lca, but may be others (e.g. gym-aloha, gym-lowcostrobot, gym-drake-lca)
  • constructs a Gym interface for simulation and runs it
  • Enables the sim to be teleoperated by keyboard or (incomplete) leader arm (note: may record multiple episodes or discard bad episodes)
  • (optional) uploads the new dataset to the Hugging Face dataset hub
  • Re-runs the recorded dataset in simulation (i.e., replays actions at correct timestamps).

How it was tested

The default arguments will run 2 episodes of 20 steps each in the PickPlaceCube-v0 environment from gym_drake_lca module and then upload them to HF dataset hub.

Examples:

This script can be tested with poetry using command:

poetry run python3 lerobot/scripts/record_and_replay_teleop_dataset.py

How to checkout & try? (for the reviewer)

Open your browser to: http://localhost:7000/ for visualization.

run:

python3 lerobot/scripts/record_and_replay_teleop_dataset.py

Control with arrows, to end recording hit END.
Look in 'data_traces/{1,2,3}/' for videos

note, the transparent black cube is the pick and place target:
Screenshot from 2024-09-18 15-44-31

The following videos should result:

  1. In "data_traces/1//videos" recordings of the teleop session with a red block.
    Screenshot from 2024-09-18 15-34-10

  2. In "data_traces/2//videos" recordings of a replay of session 1 (from disk) with a green block.
    Screenshot from 2024-09-18 15-34-03

  3. In "data_traces/3//videos" recordings of a replay of session 1 (from program memory) with a blue block.
    Screenshot from 2024-09-18 15-33-57

The example uses the low-cost robot arm single-cube examples
Other robot deployments, e.g. env-aloha, pack all image observations into one observation dict item. In consequence, this example can't be used directly, but can be used as boilerplate for those scripts.

@Cadene Cadene self-requested a review September 12, 2024 22:52
@Cadene Cadene self-assigned this Sep 12, 2024
Copy link
Collaborator

@Cadene Cadene left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for this PR :)

examples/8_record_dataset_from_gym.py Outdated Show resolved Hide resolved
examples/8_record_dataset_from_gym.py Outdated Show resolved Hide resolved
@samzapo
Copy link
Author

samzapo commented Sep 13, 2024

Closing to re-purpose as a teleoperation in sim script. Will re-open once new feature is in-place.

@samzapo samzapo closed this Sep 13, 2024
@samzapo samzapo changed the title Adds an example script for recording training data from a Gym simulation Adds a script for recording training data from a Gym simulation and replaying in sim Sep 18, 2024
@samzapo samzapo reopened this Sep 18, 2024
@samzapo
Copy link
Author

samzapo commented Oct 2, 2024

leader arm tele-op now supported:

PXL_20241002_155344954.2.mp4

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.

3 participants