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

[samzapo]

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:

The following videos should result:

1. In "data_traces/1//videos" recordings of the teleop session with a red block.
   

2. In "data_traces/2//videos" recordings of a replay of session 1 (from disk) with a green block.
   

3. In "data_traces/3//videos" recordings of a replay of session 1 (from program memory) with a blue block.
   

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]

Thanks for this PR :)

[samzapo]

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

[samzapo]

leader arm tele-op now supported:

PXL_20241002_155344954.2.mp4