Refactor: Pose randomization (#340)

## Summary
This MR moves pose randomization-on-reset into the objects, rather than
it living in the tasks.

## Detailed description
- Previously, the randomization-on-reset was handled by tasks. This is
wrong, basically, different environments want different randomization,
for the same tasks. So randomization should be specified
per-environment, not per-task.
- This MR allows `object.set_initial_pose()` to take a pose range, which
specifies randomization/
- Behind the scenes this creates a randomization event per object.

Author: alexmillane

docs/pages/example_workflows/locomanipulation/step_1_environment_setup.rst

@@ -26,7 +26,7 @@ Environment Description
                from isaaclab_arena.environments.isaaclab_arena_environment import IsaacLabArenaEnvironment
                from isaaclab_arena.scene.scene import Scene
                from isaaclab_arena.tasks.g1_locomanip_pick_and_place_task import G1LocomanipPickAndPlaceTask
-               from isaaclab_arena.utils.pose import Pose
+               from isaaclab_arena.utils.pose import Pose, PoseRange
 
                background = self.asset_registry.get_asset_by_name("galileo_locomanip")()
                pick_up_object = self.asset_registry.get_asset_by_name(args_cli.object)()
@@ -35,10 +35,13 @@ Environment Description
 
                teleop_device = self.device_registry.get_device_by_name(args_cli.teleop_device)()
 
+               XY_RANGE_M = 0.025
                pick_up_object.set_initial_pose(
-                   Pose(
-                       position_xyz=(0.5785, 0.18, 0.0707),
-                       rotation_wxyz=(0.0, 0.0, 1.0, 0.0),
+                   PoseRange(
+                       position_xyz_min=(0.5785 - XY_RANGE_M, 0.18 - XY_RANGE_M, 0.0707),
+                       position_xyz_max=(0.5785 + XY_RANGE_M, 0.18 + XY_RANGE_M, 0.0707),
+                       rpy_min=(0.0, 0.0, 0.0),
+                       rpy_max=(0.0, 0.0, 0.0),
                    )
                )
                blue_sorting_bin.set_initial_pose(
@@ -86,10 +89,13 @@ See :doc:`../../concepts/concept_assets_design` for details on asset architectur
 
 .. code-block:: python
 
+   XY_RANGE_M = 0.025
    pick_up_object.set_initial_pose(
-       Pose(
-           position_xyz=(0.5785, 0.18, 0.0707),
-           rotation_wxyz=(0.0, 0.0, 1.0, 0.0),
+       PoseRange(
+           position_xyz_min=(0.5785 - XY_RANGE_M, 0.18 - XY_RANGE_M, 0.0707),
+           position_xyz_max=(0.5785 + XY_RANGE_M, 0.18 + XY_RANGE_M, 0.0707),
+           rpy_min=(0.0, 0.0, 0.0),
+           rpy_max=(0.0, 0.0, 0.0),
        )
    )
    blue_sorting_bin.set_initial_pose(

isaaclab_arena/assets/object.py

@@ -8,11 +8,12 @@
 from isaaclab.assets import ArticulationCfg, AssetBaseCfg, RigidObjectCfg
 from isaaclab.managers import EventTermCfg, SceneEntityCfg
 from isaaclab.sim.spawners.from_files.from_files_cfg import UsdFileCfg
+from isaaclab_tasks.manager_based.manipulation.stack.mdp.franka_stack_events import randomize_object_pose
 
 from isaaclab_arena.assets.object_base import ObjectBase, ObjectType
 from isaaclab_arena.assets.object_utils import detect_object_type
 from isaaclab_arena.terms.events import set_object_pose
-from isaaclab_arena.utils.pose import Pose
+from isaaclab_arena.utils.pose import Pose, PoseRange
 from isaaclab_arena.utils.usd_helpers import has_light, open_stage
 
 
@@ -49,12 +50,19 @@ def __init__(
         self.object_cfg = self._init_object_cfg()
         self.event_cfg = self._init_event_cfg()
 
-    def set_initial_pose(self, pose: Pose) -> None:
+    def set_initial_pose(self, pose: Pose | PoseRange) -> None:
+        """Set the initial pose of the object.
+
+        Args:
+            pose: The pose to set. Can be a single pose or a pose range.
+                  In the case of a PoseRange, the object will be reset
+                  to a random pose within the range on environment reset.
+        """
         self.initial_pose = pose
         self.object_cfg = self._add_initial_pose_to_cfg(self.object_cfg)
         self.event_cfg = self._update_initial_pose_event_cfg(self.event_cfg)
 
-    def get_initial_pose(self) -> Pose | None:
+    def get_initial_pose(self) -> Pose | PoseRange | None:
         return self.initial_pose
 
     def is_initial_pose_set(self) -> bool:
@@ -131,8 +139,12 @@ def _add_initial_pose_to_cfg(
     ) -> RigidObjectCfg | ArticulationCfg | AssetBaseCfg:
         # Optionally specify initial pose
         if self.initial_pose is not None:
-            object_cfg.init_state.pos = self.initial_pose.position_xyz
-            object_cfg.init_state.rot = self.initial_pose.rotation_wxyz
+            if isinstance(self.initial_pose, Pose):
+                initial_pose = self.initial_pose
+            elif isinstance(self.initial_pose, PoseRange):
+                initial_pose = self.initial_pose.get_midpoint()
+            object_cfg.init_state.pos = initial_pose.position_xyz
+            object_cfg.init_state.rot = initial_pose.rotation_wxyz
         return object_cfg
 
     def _requires_reset_pose_event(self) -> bool:
@@ -144,24 +156,55 @@ def _requires_reset_pose_event(self) -> bool:
 
     def _init_event_cfg(self) -> EventTermCfg | None:
         if self._requires_reset_pose_event():
-            return EventTermCfg(
-                func=set_object_pose,
-                mode="reset",
-                params={
-                    "pose": self.initial_pose,
-                    "asset_cfg": SceneEntityCfg(self.name),
-                },
-            )
+            # Two possible event types:
+            # - initial pose is a Pose - reset to a single pose
+            # - initial pose is a PoseRange - reset to a random pose within the range
+            if isinstance(self.initial_pose, Pose):
+                return EventTermCfg(
+                    func=set_object_pose,
+                    mode="reset",
+                    params={
+                        "pose": self.initial_pose,
+                        "asset_cfg": SceneEntityCfg(self.name),
+                    },
+                )
+            elif isinstance(self.initial_pose, PoseRange):
+                return EventTermCfg(
+                    func=randomize_object_pose,
+                    mode="reset",
+                    params={
+                        "pose_range": self.initial_pose.to_dict(),
+                        "asset_cfgs": [SceneEntityCfg(self.name)],
+                    },
+                )
+            else:
+                raise ValueError(f"Initial pose {self.initial_pose} is not a Pose or PoseRange")
         else:
             return None
 
+    def _needs_reinit_of_event_cfg(self):
+        # If there is no event cfg, needs to be reinitialized
+        if self.event_cfg is None:
+            return True
+        # Here we check if the event cfg is for the correct pose type.
+        # If not, needs to be reinitialized.
+        if (isinstance(self.initial_pose, Pose) and ("pose" not in self.event_cfg.params)) or (
+            isinstance(self.initial_pose, PoseRange) and ("pose_range" not in self.event_cfg.params)
+        ):
+            return True
+        return False
+
     def _update_initial_pose_event_cfg(self, event_cfg: EventTermCfg | None) -> EventTermCfg | None:
         if self._requires_reset_pose_event():
             # Create an event cfg if one does not yet exist
-            if event_cfg is None:
+            if self._needs_reinit_of_event_cfg():
                 event_cfg = self._init_event_cfg()
-            # Add the initial pose to the event cfg
-            event_cfg.params["pose"] = self.initial_pose
+            if isinstance(self.initial_pose, Pose):
+                event_cfg.params["pose"] = self.initial_pose
+            elif isinstance(self.initial_pose, PoseRange):
+                event_cfg.params["pose_range"] = self.initial_pose.to_dict()
+            else:
+                raise ValueError(f"Initial pose {self.initial_pose} is not a Pose or PoseRange")
         else:
             event_cfg = None
         return event_cfg

isaaclab_arena/tasks/g1_locomanip_pick_and_place_task.py

@@ -4,16 +4,13 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import numpy as np
-import torch
 from dataclasses import MISSING
 
 import isaaclab.envs.mdp as mdp_isaac_lab
 from isaaclab.envs.common import ViewerCfg
 from isaaclab.envs.mimic_env_cfg import MimicEnvCfg, SubTaskConfig
-from isaaclab.managers import EventTermCfg, SceneEntityCfg, TerminationTermCfg
+from isaaclab.managers import SceneEntityCfg, TerminationTermCfg
 from isaaclab.utils import configclass
-from isaaclab.utils.math import euler_xyz_from_quat
-from isaaclab_tasks.manager_based.manipulation.stack.mdp import franka_stack_events
 
 from isaaclab_arena.assets.asset import Asset
 from isaaclab_arena.embodiments.common.arm_mode import ArmMode
@@ -43,6 +40,7 @@ def __init__(
             if task_description is None
             else task_description
         )
+        self.events_cfg = None
 
     def get_scene_cfg(self):
         pass
@@ -71,7 +69,7 @@ def get_termination_cfg(self):
         )
 
     def get_events_cfg(self):
-        return EventsCfg(pick_up_object=self.pick_up_object)
+        return self.events_cfg
 
     def get_mimic_env_cfg(self, arm_mode: ArmMode):
         return G1LocomanipPickPlaceMimicEnvCfg()
@@ -97,39 +95,6 @@ class TerminationsCfg:
     object_dropped: TerminationTermCfg = MISSING
 
 
-@configclass
-class EventsCfg:
-    """Configuration for Pick and Place."""
-
-    reset_pick_up_object_pose: EventTermCfg = MISSING
-
-    def __init__(self, pick_up_object: Asset):
-        initial_pose = pick_up_object.get_initial_pose()
-        if initial_pose is not None:
-            roll, pitch, yaw = euler_xyz_from_quat(torch.tensor(initial_pose.rotation_wxyz).reshape(1, 4))
-            self.reset_pick_up_object_pose = EventTermCfg(
-                func=franka_stack_events.randomize_object_pose,
-                mode="reset",
-                params={
-                    "pose_range": {
-                        "x": (initial_pose.position_xyz[0] - 0.025, initial_pose.position_xyz[0] + 0.025),
-                        "y": (initial_pose.position_xyz[1] - 0.025, initial_pose.position_xyz[1] + 0.025),
-                        "z": (initial_pose.position_xyz[2], initial_pose.position_xyz[2]),
-                        "roll": (roll, roll),
-                        "pitch": (pitch, pitch),
-                        "yaw": (yaw, yaw),
-                    },
-                    "asset_cfgs": [SceneEntityCfg(pick_up_object.name)],
-                },
-            )
-        else:
-            print(
-                f"Pick up object {pick_up_object.name} has no initial pose. Not setting reset pick up object pose"
-                " event."
-            )
-            self.reset_pick_up_object_pose = None
-
-
 @configclass
 class G1LocomanipPickPlaceMimicEnvCfg(MimicEnvCfg):
     """

isaaclab_arena/tasks/lift_object_task.py

@@ -8,7 +8,7 @@
 
 import isaaclab.envs.mdp as mdp_isaac_lab
 from isaaclab.envs.common import ViewerCfg
-from isaaclab.managers import CommandTermCfg, EventTermCfg
+from isaaclab.managers import CommandTermCfg
 from isaaclab.managers import ObservationGroupCfg as ObsGroup
 from isaaclab.managers import ObservationTermCfg as ObsTerm
 from isaaclab.managers import RewardTermCfg, SceneEntityCfg, TerminationTermCfg
@@ -38,7 +38,7 @@ def __init__(
         self.background_scene = background_scene
         self.minimum_height_to_lift = minimum_height_to_lift
         self.scene_config = None
-        self.events_cfg = LiftObjectEventsCfg(lift_object=self.lift_object, reset_pose_range=reset_pose_range.to_dict())
+        self.events_cfg = None
         self.termination_cfg = self.make_termination_cfg()
 
     def get_scene_cfg(self):
@@ -73,28 +73,6 @@ def get_viewer_cfg(self) -> ViewerCfg:
         )
 
 
-@configclass
-class LiftObjectEventsCfg:
-    """Configuration for Lift Object."""
-
-    reset_lift_object_pose: EventTermCfg = MISSING
-
-    def __init__(self, lift_object: Asset, reset_pose_range: dict[str, tuple[float, float]]):
-        self.reset_lift_object_pose = EventTermCfg(
-            func=mdp_isaac_lab.reset_root_state_uniform,
-            mode="reset",
-            params={
-                "pose_range": {
-                    "x": reset_pose_range["x"],
-                    "y": reset_pose_range["y"],
-                    "z": reset_pose_range["z"],
-                },
-                "velocity_range": {},
-                "asset_cfg": SceneEntityCfg(lift_object.name),
-            },
-        )
-
-
 @configclass
 class LiftObjectTerminationsCfg:
     """Termination terms for the Lift Object task."""
@@ -111,14 +89,12 @@ def __init__(
         embodiment: EmbodimentBase,
         minimum_height_to_lift: float = 0.04,
         episode_length_s: float = 5.0,
-        reset_pose_range: PoseRange = PoseRange(),
     ):
         super().__init__(
             lift_object=lift_object,
             background_scene=background_scene,
             minimum_height_to_lift=minimum_height_to_lift,
             episode_length_s=episode_length_s,
-            reset_pose_range=reset_pose_range,
         )
         self.embodiment = embodiment
         self.observation_cfg = LiftObjectObservationsCfg(
@@ -180,6 +156,8 @@ class LiftObjectCommandsCfg:
 
     def __init__(self, asset_name: str, body_name: str, lift_object: Asset):
         initial_pose = lift_object.get_initial_pose()
+        if isinstance(initial_pose, PoseRange):
+            initial_pose = initial_pose.get_midpoint()
         self.object_pose = mdp_isaac_lab.UniformPoseCommandCfg(
             asset_name=asset_name,
             body_name=body_name,

isaaclab_arena/tasks/pick_and_place_task.py

@@ -9,7 +9,7 @@
 import isaaclab.envs.mdp as mdp_isaac_lab
 from isaaclab.envs.common import ViewerCfg
 from isaaclab.envs.mimic_env_cfg import MimicEnvCfg, SubTaskConfig
-from isaaclab.managers import EventTermCfg, SceneEntityCfg, TerminationTermCfg
+from isaaclab.managers import SceneEntityCfg, TerminationTermCfg
 from isaaclab.sensors.contact_sensor.contact_sensor_cfg import ContactSensorCfg
 from isaaclab.utils import configclass
 
@@ -21,7 +21,6 @@
 from isaaclab_arena.tasks.task_base import TaskBase
 from isaaclab_arena.tasks.terminations import object_on_destination
 from isaaclab_arena.utils.cameras import get_viewer_cfg_look_at_object
-from isaaclab_arena.utils.pose import PoseRange
 
 
 class PickAndPlaceTask(TaskBase):
@@ -34,7 +33,6 @@ def __init__(
         destination_object: Asset | None = None,
         episode_length_s: float | None = None,
         task_description: str | None = None,
-        reset_pose_range: PoseRange = PoseRange(),
     ):
         super().__init__(episode_length_s=episode_length_s)
         self.pick_up_object = pick_up_object
@@ -46,9 +44,7 @@ def __init__(
                 contact_against_prim_paths=[self.destination_location.get_prim_path()],
             ),
         )
-        self.events_cfg = PickPlaceEventsCfg(
-            pick_up_object=self.pick_up_object, reset_pose_range=reset_pose_range.to_dict()
-        )
+        self.events_cfg = None
         self.termination_cfg = self.make_termination_cfg()
         self.task_description = (
             f"Pick up the {pick_up_object.name}, and place it into the {destination_location.name}"
@@ -122,28 +118,6 @@ class TerminationsCfg:
     object_dropped: TerminationTermCfg = MISSING
 
 
-@configclass
-class PickPlaceEventsCfg:
-    """Configuration for Pick Up Object."""
-
-    reset_pick_up_object_pose: EventTermCfg = MISSING
-
-    def __init__(self, pick_up_object: Asset, reset_pose_range: dict[str, tuple[float, float]]):
-        self.reset_pick_up_object_pose = EventTermCfg(
-            func=mdp_isaac_lab.reset_root_state_uniform,
-            mode="reset",
-            params={
-                "pose_range": {
-                    "x": reset_pose_range["x"],
-                    "y": reset_pose_range["y"],
-                    "z": reset_pose_range["z"],
-                },
-                "velocity_range": {},
-                "asset_cfg": SceneEntityCfg(pick_up_object.name),
-            },
-        )
-
-
 @configclass
 class PickPlaceMimicEnvCfg(MimicEnvCfg):
     """