diff --git a/docs/source/_static/publications/omnireset/drawer_assembly_training_curve.jpg b/docs/source/_static/publications/omnireset/drawer_assembly_training_curve.jpg deleted file mode 100644 index 13d2c24..0000000 Binary files a/docs/source/_static/publications/omnireset/drawer_assembly_training_curve.jpg and /dev/null differ diff --git a/docs/source/_static/publications/omnireset/peg_training_curve.jpg b/docs/source/_static/publications/omnireset/peg_training_curve.jpg deleted file mode 100644 index 7ba4777..0000000 Binary files a/docs/source/_static/publications/omnireset/peg_training_curve.jpg and /dev/null differ diff --git a/docs/source/_static/publications/omnireset/success_rate_over_steps.jpg b/docs/source/_static/publications/omnireset/success_rate_over_steps.jpg new file mode 100644 index 0000000..ad7e11b Binary files /dev/null and b/docs/source/_static/publications/omnireset/success_rate_over_steps.jpg differ diff --git a/docs/source/_static/publications/omnireset/success_rate_over_wall_clock.jpg b/docs/source/_static/publications/omnireset/success_rate_over_wall_clock.jpg new file mode 100644 index 0000000..8d87ac9 Binary files /dev/null and b/docs/source/_static/publications/omnireset/success_rate_over_wall_clock.jpg differ diff --git a/docs/source/_static/publications/omnireset/twisting_training_curve.jpg b/docs/source/_static/publications/omnireset/twisting_training_curve.jpg deleted file mode 100644 index 26e83a1..0000000 Binary files a/docs/source/_static/publications/omnireset/twisting_training_curve.jpg and /dev/null differ diff --git a/docs/source/_static/tasks/manipulation/omnireset_cube_stack.jpg b/docs/source/_static/tasks/manipulation/omnireset_cube_stack.jpg new file mode 100644 index 0000000..b9f3d7e Binary files /dev/null and b/docs/source/_static/tasks/manipulation/omnireset_cube_stack.jpg differ diff --git a/docs/source/_static/tasks/manipulation/omnireset_cupcake_on_plate.jpg b/docs/source/_static/tasks/manipulation/omnireset_cupcake_on_plate.jpg new file mode 100644 index 0000000..2f0d59a Binary files /dev/null and b/docs/source/_static/tasks/manipulation/omnireset_cupcake_on_plate.jpg differ diff --git a/docs/source/_static/tasks/manipulation/omnireset_rectangle_reorientation.jpg b/docs/source/_static/tasks/manipulation/omnireset_rectangle_reorientation.jpg new file mode 100644 index 0000000..97cf5d1 Binary files /dev/null and b/docs/source/_static/tasks/manipulation/omnireset_rectangle_reorientation.jpg differ diff --git a/docs/source/overview/uw_environments.rst b/docs/source/overview/uw_environments.rst index ba24670..5637cf2 100644 --- a/docs/source/overview/uw_environments.rst +++ b/docs/source/overview/uw_environments.rst @@ -57,6 +57,12 @@ Environments based on fixed-arm manipulation tasks. +--------------------------------+------------------------------------------------+------------------------------------------------------------------------------+ | |omnireset-ur5e-peg-insert| | |omnireset-ur5e-peg-insert-link| | Insert a square peg into the square peg hole | +--------------------------------+------------------------------------------------+------------------------------------------------------------------------------+ + | |omnireset-ur5e-rectangle| | |omnireset-ur5e-rectangle-link| | Orient a rectangle into desired location during a wall | + +--------------------------------+------------------------------------------------+------------------------------------------------------------------------------+ + | |omnireset-ur5e-cupcake| | |omnireset-ur5e-cupcake-link| | Place a cupcake into desired location on a plate | + +--------------------------------+------------------------------------------------+------------------------------------------------------------------------------+ + | |omnireset-ur5e-cube-stack| | |omnireset-ur5e-cube-stack-link| | Stack one cube on top of another cube in a desired orientation. | + +--------------------------------+------------------------------------------------+------------------------------------------------------------------------------+ .. |track-goal-ur5| image:: ../_static/tasks/manipulation/ur5_track_goal.jpg .. |track-goal-tycho| image:: ../_static/tasks/manipulation/tycho_track_goal.jpg @@ -67,6 +73,9 @@ Environments based on fixed-arm manipulation tasks. .. |omnireset-ur5e-drawer| image:: ../_static/tasks/manipulation/omnireset_drawer_assemble.jpg .. |omnireset-ur5e-fbleg| image:: ../_static/tasks/manipulation/omnireset_fbleg_screw.jpg .. |omnireset-ur5e-peg-insert| image:: ../_static/tasks/manipulation/omnireset_peg_insert.jpg +.. |omnireset-ur5e-rectangle| image:: ../_static/tasks/manipulation/omnireset_rectangle_reorientation.jpg +.. |omnireset-ur5e-cupcake| image:: ../_static/tasks/manipulation/omnireset_cupcake_on_plate.jpg +.. |omnireset-ur5e-cube-stack| image:: ../_static/tasks/manipulation/omnireset_cube_stack.jpg .. |track-goal-ur5-link| replace:: `UW-Track-Goal-Ur5-v0 `__ .. |track-goal-tycho-link| replace:: `UW-Track-Goal-Tycho-v0 `__ @@ -74,9 +83,12 @@ Environments based on fixed-arm manipulation tasks. .. |ext-nut-thread-franka-link| replace:: `UW-Nut-Thread-Franka-v0 `__ .. |ext-gear-mesh-franka-link| replace:: `UW-Gear-Mesh-Franka-v0 `__ .. |ext-peg-insert-franka-link| replace:: `UW-Peg-Insert-Franka-v0 `__ -.. |omnireset-ur5e-drawer-link| replace:: `OmniReset-Ur5eRobotiq2f85-RelJointPos-State-v0 `__ -.. |omnireset-ur5e-fbleg-link| replace:: `OmniReset-Ur5eRobotiq2f85-RelJointPos-State-v0 `__ -.. |omnireset-ur5e-peg-insert-link| replace:: `OmniReset-Ur5eRobotiq2f85-RelJointPos-State-v0 `__ +.. |omnireset-ur5e-drawer-link| replace:: `OmniReset-Ur5eRobotiq2f85-RelCartesianOSC-State-v0 `__ +.. |omnireset-ur5e-fbleg-link| replace:: `OmniReset-Ur5eRobotiq2f85-RelCartesianOSC-State-v0 `__ +.. |omnireset-ur5e-peg-insert-link| replace:: `OmniReset-Ur5eRobotiq2f85-RelCartesianOSC-State-v0 `__ +.. |omnireset-ur5e-rectangle-link| replace:: `OmniReset-Ur5eRobotiq2f85-RelCartesianOSC-State-v0 `__ +.. |omnireset-ur5e-cupcake-link| replace:: `OmniReset-Ur5eRobotiq2f85-RelCartesianOSC-State-v0 `__ +.. |omnireset-ur5e-cube-stack-link| replace:: `OmniReset-Ur5eRobotiq2f85-RelCartesianOSC-State-v0 `__ Locomotion diff --git a/docs/source/publications/omnireset/index.rst b/docs/source/publications/omnireset/index.rst index ee43885..884aae6 100644 --- a/docs/source/publications/omnireset/index.rst +++ b/docs/source/publications/omnireset/index.rst @@ -6,7 +6,7 @@ OmniReset .. important:: **Pre-trained RL Checkpoints Available!** - We provide trained RL checkpoints for all three tasks: **Drawer Assembly**, **Leg Twisting**, and **Peg Insertion**. + We provide trained RL checkpoints for all six tasks: **Drawer Assembly**, **Leg Twisting**, **Peg Insertion**, **Rectangle Reorientation on Wall**, **Cupcake on Plate**, and **Cube Stacking**. Download the checkpoints and evaluate them immediately! See the :doc:`instruction` guide for download links and evaluation instructions. @@ -35,6 +35,27 @@ OmniReset

Peg Insertion

+
+ +

Rectangle Reorientation on Wall

+
+
+ +

Cupcake on Plate

+
+
+ +

Cube Stacking

+
.. note:: diff --git a/docs/source/publications/omnireset/instruction.rst b/docs/source/publications/omnireset/instruction.rst index 53c7d7c..b999582 100644 --- a/docs/source/publications/omnireset/instruction.rst +++ b/docs/source/publications/omnireset/instruction.rst @@ -11,8 +11,11 @@ Choose your path: :ref:`evaluate our pre-trained checkpoints None: super().reset(env_ids) @@ -138,6 +140,11 @@ def reset(self, env_ids: torch.Tensor | None = None) -> None: object_asset.initial_pos[env_ids] = object_asset.data.root_pos_w[env_ids].clone() object_asset.initial_quat[env_ids] = object_asset.data.root_quat_w[env_ids].clone() + if env_ids is None: + self.stability_counter.zero_() + else: + self.stability_counter[env_ids] = 0 + def __call__( self, env: ManagerBasedEnv, @@ -146,6 +153,7 @@ def __call__( collision_analyzer_cfg: CollisionAnalyzerCfg, max_pos_deviation: float = 0.05, pos_z_threshold: float = 0.05, + consecutive_stability_steps: int = 5, ) -> torch.Tensor: # Get object and gripper from scene object_asset = env.scene[self.object_cfg.name] @@ -159,6 +167,26 @@ def __call__( dim=1 ) + # Check if asset velocities are small + current_step_stable = torch.ones(env.num_envs, device=env.device, dtype=torch.bool) + # Check gripper (articulation) velocities + current_step_stable &= gripper_asset.data.joint_vel.abs().sum(dim=1) < 5.0 + # Check object (rigid object) velocities + if isinstance(object_asset, RigidObject): + current_step_stable &= object_asset.data.body_lin_vel_w.abs().sum(dim=2).sum(dim=1) < 0.05 + current_step_stable &= object_asset.data.body_ang_vel_w.abs().sum(dim=2).sum(dim=1) < 1.0 + elif isinstance(object_asset, RigidObjectCollection): + current_step_stable &= object_asset.data.object_lin_vel_w.abs().sum(dim=2).sum(dim=1) < 0.05 + current_step_stable &= object_asset.data.object_ang_vel_w.abs().sum(dim=2).sum(dim=1) < 1.0 + + self.stability_counter = torch.where( + current_step_stable, + self.stability_counter + 1, # Increment counter if stable + torch.zeros_like(self.stability_counter), # Reset counter if not stable + ) + + stability_reached = self.stability_counter >= self.consecutive_stability_steps + # Skip if position or quaternion is NaN pos_is_nan = torch.isnan(object_asset.data.root_pos_w).any(dim=1) quat_is_nan = torch.isnan(object_asset.data.root_quat_w).any(dim=1) @@ -177,7 +205,12 @@ def __call__( collision_free = self.collision_analyzer(env, all_env_ids) grasp_success = ( - (~abnormal_gripper_state) & (~excessive_pose_deviation) & pos_above_ground & collision_free & time_out + (~abnormal_gripper_state) + & stability_reached + & (~excessive_pose_deviation) + & pos_above_ground + & collision_free + & time_out ) return grasp_success @@ -270,10 +303,10 @@ def __call__( if isinstance(asset, Articulation): current_step_stable &= asset.data.joint_vel.abs().sum(dim=1) < 5.0 elif isinstance(asset, RigidObject): - current_step_stable &= asset.data.body_lin_vel_w.abs().sum(dim=2).sum(dim=1) < 0.05 + current_step_stable &= asset.data.body_lin_vel_w.abs().sum(dim=2).sum(dim=1) < 0.1 current_step_stable &= asset.data.body_ang_vel_w.abs().sum(dim=2).sum(dim=1) < 1.0 elif isinstance(asset, RigidObjectCollection): - current_step_stable &= asset.data.object_lin_vel_w.abs().sum(dim=2).sum(dim=1) < 0.05 + current_step_stable &= asset.data.object_lin_vel_w.abs().sum(dim=2).sum(dim=1) < 0.1 current_step_stable &= asset.data.object_ang_vel_w.abs().sum(dim=2).sum(dim=1) < 1.0 self.stability_counter = torch.where(