Skip to content

Additional nit changes to make OmniReset docs more clear #19

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.

Sign up for GitHub

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

Merged
patrickhaoy merged 1 commit into from
Jan 4, 2026
Merged

Additional nit changes to make OmniReset docs more clear #19

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
88 changes: 56 additions & 32 deletions docs/source/publications/omnireset/index.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -117,53 +117,53 @@ Download our pretrained checkpoint and run evaluation.
env.scene.insertive_object=rectangle \
env.scene.receptive_object=wall

.. tab-item:: Cupcake on Plate
.. tab-item:: Cube Stacking

.. raw:: html

<div style="text-align: center; margin-bottom: 20px;">
<video width="400" height="300" controls>
<source src="https://s3.us-west-004.backblazeb2.com/uwlab-assets/Media/OmniReset/cupcake.mp4" type="video/mp4">
<source src="https://s3.us-west-004.backblazeb2.com/uwlab-assets/Media/OmniReset/cube.mp4" type="video/mp4">
Your browser does not support the video tag.
</video>
</div>

.. code:: bash

# Download checkpoint
wget https://s3.us-west-004.backblazeb2.com/uwlab-assets/Policies/OmniReset/cupcake_state_rl_expert.pt
wget https://s3.us-west-004.backblazeb2.com/uwlab-assets/Policies/OmniReset/cube_state_rl_expert.pt

# Run evaluation
python scripts/reinforcement_learning/rsl_rl/play.py \
--task OmniReset-Ur5eRobotiq2f85-RelCartesianOSC-State-Play-v0 \
--num_envs 1 \
--checkpoint cupcake_state_rl_expert.pt \
env.scene.insertive_object=cupcake \
env.scene.receptive_object=plate
--checkpoint cube_state_rl_expert.pt \
env.scene.insertive_object=cube \
env.scene.receptive_object=cube

.. tab-item:: Cube Stacking
.. tab-item:: Cupcake on Plate

.. raw:: html

<div style="text-align: center; margin-bottom: 20px;">
<video width="400" height="300" controls>
<source src="https://s3.us-west-004.backblazeb2.com/uwlab-assets/Media/OmniReset/cube.mp4" type="video/mp4">
<source src="https://s3.us-west-004.backblazeb2.com/uwlab-assets/Media/OmniReset/cupcake.mp4" type="video/mp4">
Your browser does not support the video tag.
</video>
</div>

.. code:: bash

# Download checkpoint
wget https://s3.us-west-004.backblazeb2.com/uwlab-assets/Policies/OmniReset/cube_state_rl_expert.pt
wget https://s3.us-west-004.backblazeb2.com/uwlab-assets/Policies/OmniReset/cupcake_state_rl_expert.pt

# Run evaluation
python scripts/reinforcement_learning/rsl_rl/play.py \
--task OmniReset-Ur5eRobotiq2f85-RelCartesianOSC-State-Play-v0 \
--num_envs 1 \
--checkpoint cube_state_rl_expert.pt \
env.scene.insertive_object=cube \
env.scene.receptive_object=cube
--checkpoint cupcake_state_rl_expert.pt \
env.scene.insertive_object=cupcake \
env.scene.receptive_object=plate

----

Expand All @@ -174,9 +174,9 @@ Reproduce Our Training

Reproduce our training results from scratch.

.. important::
.. tip::

Before running reset state generation scripts (step 3), make sure ``base_path`` and ``base_paths`` in ``reset_states_cfg.py`` are set appropriately.
**Want to try it quickly?** Start with **Cube Stacking** or **Peg Insertion**. They have the fastest reset state collection times and converge within ~8 hours on 4×L40S GPUs.

.. tab-set::

Expand All @@ -196,6 +196,10 @@ Reproduce our training results from scratch.

**Step 3: Generate Reset State Datasets** (~1 min to 1 hour depending on the reset)

.. important::

Before running, make sure ``base_path`` and ``base_paths`` in ``reset_states_cfg.py`` are set appropriately.

.. code:: bash

# Object Anywhere, End-Effector Anywhere (Reaching)
Expand Down Expand Up @@ -242,6 +246,10 @@ Reproduce our training results from scratch.

**Step 3: Generate Reset State Datasets** (~1 min to 1 hour depending on the reset)

.. important::

Before running, make sure ``base_path`` and ``base_paths`` in ``reset_states_cfg.py`` are set appropriately.

.. code:: bash

# Object Anywhere, End-Effector Anywhere (Reaching)
Expand Down Expand Up @@ -288,6 +296,10 @@ Reproduce our training results from scratch.

**Step 3: Generate Reset State Datasets** (~1 min to 1 hour depending on the reset)

.. important::

Before running, make sure ``base_path`` and ``base_paths`` in ``reset_states_cfg.py`` are set appropriately.

.. code:: bash

# Object Anywhere, End-Effector Anywhere (Reaching)
Expand Down Expand Up @@ -334,6 +346,10 @@ Reproduce our training results from scratch.

**Step 3: Generate Reset State Datasets** (~1 min to 1 hour depending on the reset)

.. important::

Before running, make sure ``base_path`` and ``base_paths`` in ``reset_states_cfg.py`` are set appropriately.

.. code:: bash

# Object Anywhere, End-Effector Anywhere (Reaching)
Expand Down Expand Up @@ -364,35 +380,39 @@ Reproduce our training results from scratch.
env.scene.insertive_object=rectangle \
env.scene.receptive_object=wall

.. tab-item:: Cupcake on Plate
.. tab-item:: Cube Stacking

**Step 1: Collect Partial Assemblies** (~30 seconds)

.. code:: bash

python scripts_v2/tools/record_partial_assemblies.py --task OmniReset-PartialAssemblies-v0 --num_envs 10 --num_trajectories 10 --dataset_dir ./partial_assembly_datasets --headless env.scene.insertive_object=cupcake env.scene.receptive_object=plate
python scripts_v2/tools/record_partial_assemblies.py --task OmniReset-PartialAssemblies-v0 --num_envs 10 --num_trajectories 10 --dataset_dir ./partial_assembly_datasets --headless env.scene.insertive_object=cube env.scene.receptive_object=cube

**Step 2: Sample Grasp Poses** (~1 minute)

.. code:: bash

python scripts_v2/tools/record_grasps.py --task OmniReset-Robotiq2f85-GraspSampling-v0 --num_envs 8192 --num_grasps 1000 --dataset_dir ./grasp_datasets --headless env.scene.object=cupcake
python scripts_v2/tools/record_grasps.py --task OmniReset-Robotiq2f85-GraspSampling-v0 --num_envs 8192 --num_grasps 1000 --dataset_dir ./grasp_datasets --headless env.scene.object=cube

**Step 3: Generate Reset State Datasets** (~1 min to 1 hour depending on the reset)

.. important::

Before running, make sure ``base_path`` and ``base_paths`` in ``reset_states_cfg.py`` are set appropriately.

.. code:: bash

# Object Anywhere, End-Effector Anywhere (Reaching)
python scripts_v2/tools/record_reset_states.py --task OmniReset-UR5eRobotiq2f85-ObjectAnywhereEEAnywhere-v0 --num_envs 4096 --num_reset_states 10000 --headless --dataset_dir ./reset_state_datasets/ObjectAnywhereEEAnywhere env.scene.insertive_object=cupcake env.scene.receptive_object=plate
python scripts_v2/tools/record_reset_states.py --task OmniReset-UR5eRobotiq2f85-ObjectAnywhereEEAnywhere-v0 --num_envs 4096 --num_reset_states 10000 --headless --dataset_dir ./reset_state_datasets/ObjectAnywhereEEAnywhere env.scene.insertive_object=cube env.scene.receptive_object=cube

# Object Resting, End-Effector Grasped (Near Object)
python scripts_v2/tools/record_reset_states.py --task OmniReset-UR5eRobotiq2f85-ObjectRestingEEGrasped-v0 --num_envs 4096 --num_reset_states 10000 --headless --dataset_dir ./reset_state_datasets/ObjectRestingEEGrasped env.scene.insertive_object=cupcake env.scene.receptive_object=plate
python scripts_v2/tools/record_reset_states.py --task OmniReset-UR5eRobotiq2f85-ObjectRestingEEGrasped-v0 --num_envs 4096 --num_reset_states 10000 --headless --dataset_dir ./reset_state_datasets/ObjectRestingEEGrasped env.scene.insertive_object=cube env.scene.receptive_object=cube

# Object Anywhere, End-Effector Grasped (Grasped)
python scripts_v2/tools/record_reset_states.py --task OmniReset-UR5eRobotiq2f85-ObjectAnywhereEEGrasped-v0 --num_envs 4096 --num_reset_states 10000 --headless --dataset_dir ./reset_state_datasets/ObjectAnywhereEEGrasped env.scene.insertive_object=cupcake env.scene.receptive_object=plate
python scripts_v2/tools/record_reset_states.py --task OmniReset-UR5eRobotiq2f85-ObjectAnywhereEEGrasped-v0 --num_envs 4096 --num_reset_states 10000 --headless --dataset_dir ./reset_state_datasets/ObjectAnywhereEEGrasped env.scene.insertive_object=cube env.scene.receptive_object=cube

# Object Partially Assembled, End-Effector Grasped (Near Goal)
python scripts_v2/tools/record_reset_states.py --task OmniReset-UR5eRobotiq2f85-ObjectPartiallyAssembledEEGrasped-v0 --num_envs 4096 --num_reset_states 10000 --headless --dataset_dir ./reset_state_datasets/ObjectPartiallyAssembledEEGrasped env.scene.insertive_object=cupcake env.scene.receptive_object=plate
python scripts_v2/tools/record_reset_states.py --task OmniReset-UR5eRobotiq2f85-ObjectPartiallyAssembledEEGrasped-v0 --num_envs 4096 --num_reset_states 10000 --headless --dataset_dir ./reset_state_datasets/ObjectPartiallyAssembledEEGrasped env.scene.insertive_object=cube env.scene.receptive_object=cube

**Step 4: Train RL Policy**

Expand All @@ -407,38 +427,42 @@ Reproduce our training results from scratch.
--logger wandb \
--headless \
--distributed \
env.scene.insertive_object=cupcake \
env.scene.receptive_object=plate
env.scene.insertive_object=cube \
env.scene.receptive_object=cube

.. tab-item:: Cube Stacking
.. tab-item:: Cupcake on Plate

**Step 1: Collect Partial Assemblies** (~30 seconds)

.. code:: bash

python scripts_v2/tools/record_partial_assemblies.py --task OmniReset-PartialAssemblies-v0 --num_envs 10 --num_trajectories 10 --dataset_dir ./partial_assembly_datasets --headless env.scene.insertive_object=cube env.scene.receptive_object=cube
python scripts_v2/tools/record_partial_assemblies.py --task OmniReset-PartialAssemblies-v0 --num_envs 10 --num_trajectories 10 --dataset_dir ./partial_assembly_datasets --headless env.scene.insertive_object=cupcake env.scene.receptive_object=plate

**Step 2: Sample Grasp Poses** (~1 minute)

.. code:: bash

python scripts_v2/tools/record_grasps.py --task OmniReset-Robotiq2f85-GraspSampling-v0 --num_envs 8192 --num_grasps 1000 --dataset_dir ./grasp_datasets --headless env.scene.object=cube
python scripts_v2/tools/record_grasps.py --task OmniReset-Robotiq2f85-GraspSampling-v0 --num_envs 8192 --num_grasps 1000 --dataset_dir ./grasp_datasets --headless env.scene.object=cupcake

**Step 3: Generate Reset State Datasets** (~1 min to 1 hour depending on the reset)

.. important::

Before running, make sure ``base_path`` and ``base_paths`` in ``reset_states_cfg.py`` are set appropriately.

.. code:: bash

# Object Anywhere, End-Effector Anywhere (Reaching)
python scripts_v2/tools/record_reset_states.py --task OmniReset-UR5eRobotiq2f85-ObjectAnywhereEEAnywhere-v0 --num_envs 4096 --num_reset_states 10000 --headless --dataset_dir ./reset_state_datasets/ObjectAnywhereEEAnywhere env.scene.insertive_object=cube env.scene.receptive_object=cube
python scripts_v2/tools/record_reset_states.py --task OmniReset-UR5eRobotiq2f85-ObjectAnywhereEEAnywhere-v0 --num_envs 4096 --num_reset_states 10000 --headless --dataset_dir ./reset_state_datasets/ObjectAnywhereEEAnywhere env.scene.insertive_object=cupcake env.scene.receptive_object=plate

# Object Resting, End-Effector Grasped (Near Object)
python scripts_v2/tools/record_reset_states.py --task OmniReset-UR5eRobotiq2f85-ObjectRestingEEGrasped-v0 --num_envs 4096 --num_reset_states 10000 --headless --dataset_dir ./reset_state_datasets/ObjectRestingEEGrasped env.scene.insertive_object=cube env.scene.receptive_object=cube
python scripts_v2/tools/record_reset_states.py --task OmniReset-UR5eRobotiq2f85-ObjectRestingEEGrasped-v0 --num_envs 4096 --num_reset_states 10000 --headless --dataset_dir ./reset_state_datasets/ObjectRestingEEGrasped env.scene.insertive_object=cupcake env.scene.receptive_object=plate

# Object Anywhere, End-Effector Grasped (Grasped)
python scripts_v2/tools/record_reset_states.py --task OmniReset-UR5eRobotiq2f85-ObjectAnywhereEEGrasped-v0 --num_envs 4096 --num_reset_states 10000 --headless --dataset_dir ./reset_state_datasets/ObjectAnywhereEEGrasped env.scene.insertive_object=cube env.scene.receptive_object=cube
python scripts_v2/tools/record_reset_states.py --task OmniReset-UR5eRobotiq2f85-ObjectAnywhereEEGrasped-v0 --num_envs 4096 --num_reset_states 10000 --headless --dataset_dir ./reset_state_datasets/ObjectAnywhereEEGrasped env.scene.insertive_object=cupcake env.scene.receptive_object=plate

# Object Partially Assembled, End-Effector Grasped (Near Goal)
python scripts_v2/tools/record_reset_states.py --task OmniReset-UR5eRobotiq2f85-ObjectPartiallyAssembledEEGrasped-v0 --num_envs 4096 --num_reset_states 10000 --headless --dataset_dir ./reset_state_datasets/ObjectPartiallyAssembledEEGrasped env.scene.insertive_object=cube env.scene.receptive_object=cube
python scripts_v2/tools/record_reset_states.py --task OmniReset-UR5eRobotiq2f85-ObjectPartiallyAssembledEEGrasped-v0 --num_envs 4096 --num_reset_states 10000 --headless --dataset_dir ./reset_state_datasets/ObjectPartiallyAssembledEEGrasped env.scene.insertive_object=cupcake env.scene.receptive_object=plate

**Step 4: Train RL Policy**

Expand All @@ -453,8 +477,8 @@ Reproduce our training results from scratch.
--logger wandb \
--headless \
--distributed \
env.scene.insertive_object=cube \
env.scene.receptive_object=cube
env.scene.insertive_object=cupcake \
env.scene.receptive_object=plate

Training Curves
^^^^^^^^^^^^^^^
Expand Down
Loading