Merge branch 'og-develop' into scene-graph-multi-agents

omnigibson/action_primitives/starter_semantic_action_primitives.py

@@ -33,7 +33,9 @@
 from omnigibson.macros import create_module_macros
 from omnigibson.objects.object_base import BaseObject
 from omnigibson.objects.usd_object import USDObject
-from omnigibson.robots import BaseRobot, Fetch, Tiago
+from omnigibson.robots import *
+from omnigibson.robots.locomotion_robot import LocomotionRobot
+from omnigibson.robots.manipulation_robot import ManipulationRobot
 from omnigibson.tasks.behavior_task import BehaviorTask
 from omnigibson.utils.control_utils import FKSolver, IKSolver
 from omnigibson.utils.grasping_planning_utils import get_grasp_poses_for_object_sticky, get_grasp_position_for_open
@@ -51,8 +53,26 @@
 
 m.DEFAULT_BODY_OFFSET_FROM_FLOOR = 0.01
 
-m.KP_LIN_VEL = 0.3
-m.KP_ANGLE_VEL = 0.2
+m.KP_LIN_VEL = {
+    Tiago: 0.3,
+    Fetch: 0.3,
+    Stretch: 0.5,
+    Turtlebot: 0.3,
+    Husky: 0.05,
+    Freight: 0.05,
+    Locobot: 0.8,
+    BehaviorRobot: 0.3,
+}
+m.KP_ANGLE_VEL = {
+    Tiago: 0.2,
+    Fetch: 0.2,
+    Stretch: 0.8,
+    Turtlebot: 0.2,
+    Husky: 0.05,
+    Freight: 0.05,
+    Locobot: 2.0,
+    BehaviorRobot: 0.2,
+}
 
 m.MAX_STEPS_FOR_SETTLING = 500
 
@@ -114,6 +134,9 @@ def __init__(self, env, robot, robot_copy, robot_copy_type="original"):
         self.robot_copy_type = robot_copy_type if robot_copy_type in robot_copy.prims.keys() else "original"
         self.disabled_collision_pairs_dict = {}
 
+        # For now, the planning context only works with Fetch and Tiago
+        assert isinstance(self.robot, (Fetch, Tiago)), "PlanningContext only works with Fetch and Tiago."
+
     def __enter__(self):
         self._assemble_robot_copy()
         self._construct_disabled_collision_pairs()
@@ -285,23 +308,18 @@ def __init__(
             StarterSemanticActionPrimitiveSet.TOGGLE_OFF: self._toggle_off,
         }
         # Validate the robot
-        assert isinstance(
-            self.robot, (Fetch, Tiago)
-        ), "StarterSemanticActionPrimitives only works with Fetch and Tiago."
-        assert isinstance(
-            self.robot.controllers["base"], (JointController, DifferentialDriveController)
-        ), "StarterSemanticActionPrimitives only works with a JointController or DifferentialDriveController at the robot base."
-        self._base_controller_is_joint = isinstance(self.robot.controllers["base"], JointController)
-        if self._base_controller_is_joint:
-            assert not self.robot.controllers[
-                "base"
-            ].use_delta_commands, (
-                "StarterSemanticActionPrimitives only works with a base JointController with absolute mode."
-            )
+        if isinstance(self.robot, LocomotionRobot):
+            assert isinstance(
+                self.robot.controllers["base"], (JointController, DifferentialDriveController)
+            ), "StarterSemanticActionPrimitives only works with a JointController or DifferentialDriveController at the robot base."
+            if self._base_controller_is_joint:
+                assert not self.robot.controllers[
+                    "base"
+                ].use_delta_commands, (
+                    "StarterSemanticActionPrimitives only works with a base JointController with absolute mode."
+                )
 
-        self.arm = self.robot.default_arm
         self.robot_model = self.robot.model_name
-        self.robot_base_mass = self.robot._links["base_link"].mass
         self.add_context = add_context
 
         self._task_relevant_objects_only = task_relevant_objects_only
@@ -312,6 +330,16 @@ def __init__(
 
         self.robot_copy = self._load_robot_copy()
 
+    @property
+    def arm(self):
+        if not isinstance(self.robot, ManipulationRobot):
+            raise ValueError("Cannot use arm for non-manipulation robot")
+        return self.robot.default_arm
+
+    @property
+    def _base_controller_is_joint(self):
+        return isinstance(self.robot.controllers["base"], JointController)
+
     def _postprocess_action(self, action):
         """Postprocesses action by applying head tracking and adding context if necessary."""
         if self._enable_head_tracking:
@@ -340,12 +368,6 @@ def _load_robot_copy(self):
         robot_copy = RobotCopy()
 
         robots_to_copy = {"original": {"robot": self.robot, "copy_path": self.robot.prim_path + "_copy"}}
-        if hasattr(self.robot, "simplified_mesh_usd_path"):
-            simplified_robot = {
-                "robot": USDObject("simplified_copy", self.robot.simplified_mesh_usd_path),
-                "copy_path": "/World/simplified_robot_copy",
-            }
-            robots_to_copy["simplified"] = simplified_robot
 
         for robot_type, rc in robots_to_copy.items():
             copy_robot = None
@@ -1414,24 +1436,17 @@ def _empty_action(self):
         """
         action = np.zeros(self.robot.action_dim)
         for name, controller in self.robot._controllers.items():
-            joint_idx = controller.dof_idx
             action_idx = self.robot.controller_action_idx[name]
-            if (
-                controller.control_type == ControlType.POSITION
-                and len(joint_idx) == len(action_idx)
-                and not controller.use_delta_commands
-            ):
-                action[action_idx] = self.robot.get_joint_positions()[joint_idx]
-            elif self.robot._controller_config[name]["name"] == "InverseKinematicsController":
-                # overwrite the goal orientation, since it is in absolute frame.
+            no_op_goal = controller.compute_no_op_goal(self.robot.get_control_dict())
+
+            if self.robot._controller_config[name]["name"] == "InverseKinematicsController":
                 assert (
                     self.robot._controller_config["arm_" + self.arm]["mode"] == "pose_absolute_ori"
                 ), "Controller must be in pose_absolute_ori mode"
-                current_quat = self.robot.get_relative_eef_orientation()
-                current_ori = T.quat2axisangle(current_quat)
-                control_idx = self.robot.controller_action_idx["arm_" + self.arm]
-                action[control_idx[3:]] = current_ori
+                # convert quaternion to axis-angle representation for control input
+                no_op_goal["target_quat"] = T.quat2axisangle(no_op_goal["target_quat"])
 
+            action[action_idx] = np.concatenate(list(no_op_goal.values()))
         return action
 
     def _reset_hand(self):
@@ -1465,10 +1480,12 @@ def _get_reset_eef_pose(self):
             return np.array([0.28493954, 0.37450749, 1.1512334]), np.array(
                 [-0.21533823, 0.05361032, -0.08631776, 0.97123871]
             )
-        else:
+        elif self.robot_model == "Fetch":
             return np.array([0.48688125, -0.12507881, 0.97888719]), np.array(
                 [0.61324748, 0.61305553, -0.35266518, 0.35173529]
             )
+        else:
+            raise ValueError(f"Unsupported robot model: {self.robot_model}")
 
     def _get_reset_joint_pos(self):
         reset_pose_fetch = np.array(
@@ -1521,7 +1538,12 @@ def _get_reset_joint_pos(self):
                 4.50000000e-02,
             ]
         )
-        return reset_pose_tiago if self.robot_model == "Tiago" else reset_pose_fetch
+        if self.robot_model == "Fetch":
+            return reset_pose_fetch
+        elif self.robot_model == "Tiago":
+            return reset_pose_tiago
+        else:
+            raise ValueError(f"Unsupported robot model: {self.robot_model}")
 
     def _navigate_to_pose(self, pose_2d):
         """
@@ -1642,11 +1664,20 @@ def _navigate_to_pose_direct(self, pose_2d, low_precision=False):
             else:
                 action = self._empty_action()
                 if self._base_controller_is_joint:
-                    direction_vec = body_target_pose[0][:2] / np.linalg.norm(body_target_pose[0][:2]) * m.KP_LIN_VEL
+                    base_action_size = self.robot.controller_action_idx["base"].size
+                    assert (
+                        base_action_size == 3
+                    ), "Currently, the action primitives only support [x, y, theta] joint controller"
+                    direction_vec = (
+                        body_target_pose[0][:2]
+                        / np.linalg.norm(body_target_pose[0][:2])
+                        * m.KP_LIN_VEL[type(self.robot)]
+                    )
                     base_action = [direction_vec[0], direction_vec[1], 0.0]
                     action[self.robot.controller_action_idx["base"]] = base_action
                 else:
-                    base_action = [m.KP_LIN_VEL, 0.0]
+                    # Diff drive controller
+                    base_action = [m.KP_LIN_VEL[type(self.robot)], 0.0]
                     action[self.robot.controller_action_idx["base"]] = base_action
                 yield self._postprocess_action(action)
 
@@ -1682,9 +1713,22 @@ def _rotate_in_place(self, end_pose, angle_threshold=m.DEFAULT_ANGLE_THRESHOLD):
             action = self._empty_action()
 
             direction = -1.0 if diff_yaw < 0.0 else 1.0
-            ang_vel = m.KP_ANGLE_VEL * direction
+            ang_vel = m.KP_ANGLE_VEL[type(self.robot)] * direction
+
+            if isinstance(self.robot, Locobot) or isinstance(self.robot, Freight):
+                # Locobot and Freight wheel joints are reversed
+                ang_vel = -ang_vel
+
+            base_action = action[self.robot.controller_action_idx["base"]]
+
+            if not self._base_controller_is_joint:
+                base_action[1] = ang_vel
+            else:
+                assert (
+                    base_action.size == 3
+                ), "Currently, the action primitives only support [x, y, theta] joint controller"
+                base_action[2] = ang_vel
 
-            base_action = [0.0, 0.0, ang_vel] if self._base_controller_is_joint else [0.0, ang_vel]
             action[self.robot.controller_action_idx["base"]] = base_action
             yield self._postprocess_action(action)
 

omnigibson/objects/object_base.py

@@ -192,6 +192,9 @@ def _post_load(self):
             self.root_prim.RemoveAPI(lazy.pxr.UsdPhysics.ArticulationRootAPI)
             self.root_prim.RemoveAPI(lazy.pxr.PhysxSchema.PhysxArticulationAPI)
 
+        if og.sim.is_playing:
+            log.warning("Cannot set articulation root API while simulation is playing!")
+
         # Potentially add articulation root APIs and also set self collisions
         root_prim = (
             None

omnigibson/robots/behavior_robot.py

@@ -316,10 +316,10 @@ def _default_controller_config(self):
         )
         return controllers
 
-    def load(self):
-        prim = super(BehaviorRobot, self).load()
+    def load(self, scene):
+        prim = super(BehaviorRobot, self).load(scene)
         for part in self.parts.values():
-            part.load()
+            part.load(scene)
         return prim
 
     def _post_load(self):
@@ -486,29 +486,30 @@ class BRPart(ABC):
     """This is the interface that all BehaviorRobot eef parts must implement."""
 
     def __init__(
-        self, name: str, parent: BehaviorRobot, prim_path: str, eef_type: str, offset_to_body: List[float]
+        self, name: str, parent: BehaviorRobot, relative_prim_path: str, eef_type: str, offset_to_body: List[float]
     ) -> None:
         """
         Create an object instance with the minimum information of class ID and rendering parameters.
 
         Args:
             name (str): unique name of this BR part
             parent (BehaviorRobot): the parent BR object
-            prim_path (str): prim path to the root link of the eef
+            relative_prim_path (str): relative prim path to the root link of the eef
             eef_type (str): type of eef. One of hand, head
             offset_to_body (List[float]): relative POSITION offset between the rz link and the eef link.
         """
         self.name = name
         self.parent = parent
-        self.prim_path = prim_path
+        self.relative_prim_path = relative_prim_path
         self.eef_type = eef_type
         self.offset_to_body = offset_to_body
 
         self.ghost_hand = None
         self._root_link = None
 
-    def load(self) -> None:
-        self._root_link = self.parent.links[self.prim_path]
+    def load(self, scene) -> None:
+        self.scene = scene
+        self._root_link = self.parent.links[self.relative_prim_path.replace("/", "")]
         # setup ghost hand
         if self.eef_type == "hand" and self.parent._use_ghost_hands:
             gh_name = f"ghost_hand_{self.name}"

omnigibson/robots/fetch.py

@@ -226,7 +226,7 @@ def _post_load(self):
 
         # Also apply a convex decomposition to the torso lift link
         torso_lift_link = self.links["torso_lift_link"]
-        assert set(torso_lift_link.collision_meshes) == {"collisions"}, "Wheel link should only have 1 collision!"
+        assert set(torso_lift_link.collision_meshes) == {"collisions"}, "torso link should only have 1 collision!"
         torso_lift_link.collision_meshes["collisions"].set_collision_approximation("convexDecomposition")
 
     @property

omnigibson/robots/freight.py

@@ -4,6 +4,9 @@
 
 from omnigibson.macros import gm
 from omnigibson.robots.two_wheel_robot import TwoWheelRobot
+from omnigibson.utils.ui_utils import create_module_logger
+
+log = create_module_logger(module_name=__name__)
 
 
 class Freight(TwoWheelRobot):
@@ -13,6 +16,19 @@ class Freight(TwoWheelRobot):
     Uses joint velocity control
     """
 
+    def _post_load(self):
+        super()._post_load()
+
+        # Set the wheels back to using sphere approximations
+        for wheel_name in ["l_wheel_link", "r_wheel_link"]:
+            log.warning(
+                "Freight wheel links are post-processed to use sphere approximation collision meshes. "
+                "Please ignore any previous errors about these collision meshes."
+            )
+            wheel_link = self.links[wheel_name]
+            assert set(wheel_link.collision_meshes) == {"collisions"}, "Wheel link should only have 1 collision!"
+            wheel_link.collision_meshes["collisions"].set_collision_approximation("boundingSphere")
+
     @property
     def wheel_radius(self):
         return 0.0613

omnigibson/robots/locomotion_robot.py

@@ -1,12 +1,11 @@
 from abc import abstractmethod
 
 import numpy as np
-from transforms3d.euler import euler2quat
-from transforms3d.quaternions import qmult, quat2mat
 
 from omnigibson.controllers import LocomotionController
 from omnigibson.robots.robot_base import BaseRobot
 from omnigibson.utils.python_utils import classproperty
+from omnigibson.utils.transform_utils import euler2quat, quat2mat, quat_multiply
 from omnigibson.utils.usd_utils import ControllableObjectViewAPI
 
 
@@ -171,7 +170,7 @@ def turn_left(self, delta=0.03):
             delta (float): delta angle to rotate the base left
         """
         quat = self.get_orientation()
-        quat = qmult((euler2quat(-delta, 0, 0)), quat)
+        quat = quat_multiply((euler2quat(-delta, 0, 0)), quat)
         self.set_orientation(quat)
 
     def turn_right(self, delta=0.03):
@@ -182,7 +181,7 @@ def turn_right(self, delta=0.03):
             delta (float): angle to rotate the base right
         """
         quat = self.get_orientation()
-        quat = qmult((euler2quat(delta, 0, 0)), quat)
+        quat = quat_multiply((euler2quat(delta, 0, 0)), quat)
         self.set_orientation(quat)
 
     @property

omnigibson/robots/stretch.py

@@ -6,6 +6,9 @@
 from omnigibson.robots.active_camera_robot import ActiveCameraRobot
 from omnigibson.robots.manipulation_robot import GraspingPoint, ManipulationRobot
 from omnigibson.robots.two_wheel_robot import TwoWheelRobot
+from omnigibson.utils.ui_utils import create_module_logger
+
+log = create_module_logger(module_name=__name__)
 
 
 class Stretch(ManipulationRobot, TwoWheelRobot, ActiveCameraRobot):
@@ -14,6 +17,19 @@ class Stretch(ManipulationRobot, TwoWheelRobot, ActiveCameraRobot):
     Reference: https://hello-robot.com/stretch-3-product
     """
 
+    def _post_load(self):
+        super()._post_load()
+
+        # Set the wheels back to using sphere approximations
+        for wheel_name in ["link_left_wheel", "link_right_wheel"]:
+            log.warning(
+                "Stretch wheel links are post-processed to use sphere approximation collision meshes. "
+                "Please ignore any previous errors about these collision meshes."
+            )
+            wheel_link = self.links[wheel_name]
+            assert set(wheel_link.collision_meshes) == {"collisions"}, "Wheel link should only have 1 collision!"
+            wheel_link.collision_meshes["collisions"].set_collision_approximation("boundingSphere")
+
     @property
     def discrete_action_list(self):
         # Not supported for this robot