Merge pull request #1073 from StanfordVL/fixes/controller

Fix various controller issues

omnigibson/action_primitives/curobo.py

@@ -1,5 +1,4 @@
 import math
-from collections.abc import Iterable
 from enum import Enum
 
 import torch as th  # MUST come before importing omni!!!
@@ -8,11 +7,11 @@
 import omnigibson.lazy as lazy
 import omnigibson.utils.transform_utils as T
 from omnigibson.macros import create_module_macros
-from omnigibson.object_states.factory import METALINK_PREFIXES
 from omnigibson.prims.rigid_prim import RigidPrim
 from omnigibson.robots.holonomic_base_robot import HolonomicBaseRobot
 from omnigibson.utils.constants import JointType
 
+
 # Gives 1 - 5% better speedup, according to https://github.com/NVlabs/curobo/discussions/245#discussioncomment-9265692
 th.backends.cudnn.benchmark = True
 th.backends.cuda.matmul.allow_tf32 = True
@@ -58,81 +57,6 @@ def create_world_mesh_collision(tensor_args, obb_cache_size=10, mesh_cache_size=
     return lazy.curobo.geom.sdf.utils.create_collision_checker(world_cfg)
 
 
-def get_obstacles(
-    reference_prim_path=None,
-    only_paths=None,
-    ignore_paths=None,
-    only_substring=None,
-    ignore_substring=None,
-):
-    """
-    Grabs world collision representation
-
-    Args:
-        reference_prim_path (None or str): If specified, the prim path defining the collision world frame
-        only_paths (None or list of str): If specified, only include these sets of prim paths in the resulting
-            collision representation
-        ignore_paths (None or list of str): If specified, exclude these sets of prim paths from the resulting
-            collision representation
-        only_substring (None or list of str): If specified, only include any prim path that includes any of these
-            substrings in the resulting collision representation
-        ignore_substring (None or list of str): If specified, exclude any prim path that includes any of these substrings
-            from the resulting collision representation
-    """
-    usd_help = lazy.curobo.util.usd_helper.UsdHelper()
-    usd_help.stage = og.sim.stage
-    return usd_help.get_obstacles_from_stage(
-        reference_prim_path=reference_prim_path,
-        only_paths=only_paths,
-        ignore_paths=ignore_paths,
-        only_substring=only_substring,
-        ignore_substring=ignore_substring,
-    ).get_collision_check_world()  # WorldConfig
-
-
-def get_obstacles_sphere_representation(
-    obstacles,
-    tensor_args,
-    n_spheres=20,
-    sphere_radius=0.001,
-):
-    """
-    Gest the collision sphere representation of obstacles @obstacles
-    Args:
-        obstacles (list of Obstacle): Obstacles whose aggregate sphere representation will be computed
-        tensor_args (TensorDeviceType): Tensor device information
-        n_spheres (int or list of int): Either per-obstacle or default number of collision spheres for representing
-            each obstacle
-        sphere_radius (float or list of float): Either per-obstacle or default radius of collision spheres for
-            representing each obstacle
-
-    Returns:
-        th.Tensor: (N, 4)-shaped tensor, where each of the N computed collision spheres are defined by (x,y,z,r),
-            where (x,y,z) is the global position and r defines the sphere radius
-    """
-    n_obstacles = len(obstacles)
-    if not isinstance(n_spheres, Iterable):
-        n_spheres = [n_spheres] * n_obstacles
-    if not isinstance(sphere_radius, Iterable):
-        sphere_radius = [sphere_radius] * n_obstacles
-
-    sph_list = []
-    for obs, n_sph, sph_radius in zip(obstacles, n_spheres, sphere_radius):
-        sph = obs.get_bounding_spheres(
-            n_sph,
-            sph_radius,
-            pre_transform_pose=lazy.curobo.types.math.Pose(
-                position=th.zeros(3), quaternion=th.tensor([1.0, 0, 0, 0])
-            ).to(tensor_args),
-            tensor_args=tensor_args,
-            fit_type=lazy.curobo.geom.sphere_fit.SphereFitType.VOXEL_VOLUME_SAMPLE_SURFACE,
-            voxelize_method="ray",
-        )
-        sph_list += [s.position + [s.radius] for s in sph]
-
-    return tensor_args.to_device(th.as_tensor(sph_list))
-
-
 class CuRoboMotionGenerator:
     """
     Class for motion generator using CuRobo backend
@@ -194,6 +118,10 @@ def __init__(
             self._tensor_args, obb_cache_size=10, mesh_cache_size=2048, max_distance=0.05
         )
 
+        usd_help = lazy.curobo.util.usd_helper.UsdHelper()
+        usd_help.stage = og.sim.stage
+        self.usd_help = usd_help
+
         self.mg = dict()
         self.ee_link = dict()
         self.additional_links = dict()
@@ -282,55 +210,47 @@ def save_visualization(self, q, file_path, emb_sel=CuRoboEmbodimentSelection.DEF
         robot_world.add_obstacle(mesh_world.mesh[0])
         robot_world.save_world_as_mesh(file_path)
 
-    def update_obstacles(self, ignore_paths=None):
+    def update_obstacles(self):
         """
         Updates internal world collision cache representation based on sim state
-
-        Args:
-            ignore_paths (None or list of str): If specified, prim path substrings that should
-                be ignored when updating obstacles
         """
-        print("Updating CuRobo world, reading w.r.t.", self.robot.prim_path)
-        ignore_paths = [] if ignore_paths is None else ignore_paths
-
-        # Ignore any visual only objects and any objects not part of the robot's current scene
-        ignore_scenes = [scene.prim_path for scene in og.sim.scenes]
-        del ignore_scenes[self.robot.scene.idx]
-        ignore_visual_only = [obj.prim_path for obj in self.robot.scene.objects if obj.visual_only]
-
-        obstacles = get_obstacles(
-            reference_prim_path=self.robot.root_link.prim_path,
-            ignore_substring=[
-                self.robot.prim_path,  # Don't include robot paths
-                "/curobo",  # Don't include curobo prim
-                "visual",  # Don't include any visuals
-                *METALINK_PREFIXES,  # Don't include any metalinks
-                *ignore_scenes,  # Don't include any scenes the robot is not in
-                *ignore_visual_only,  # Don't include any visual-only objects
-                *ignore_paths,  # Don't include any additional specified paths
-            ],
-        )
-        # All embodiment selections share the same world collision checker
-        self.mg[CuRoboEmbodimentSelection.DEFAULT].update_world(obstacles)
-        print("Synced CuRobo world from stage.")
-
-    def update_obstacles_fast(self):
-        # All embodiment selections share the same world collision checker
-        world_coll_checker = self.mg[CuRoboEmbodimentSelection.DEFAULT].world_coll_checker
-        for i, prim_path in enumerate(world_coll_checker._env_mesh_names[0]):
-            if prim_path is None:
+        obstacles = {"cuboid": None, "sphere": None, "mesh": [], "cylinder": None, "capsule": None}
+        robot_transform = T.pose_inv(T.pose2mat(self.robot.root_link.get_position_orientation()))
+
+        if og.sim.floor_plane is not None:
+            prim = og.sim.floor_plane.prim.GetChildren()[0]
+            m = lazy.curobo.util.usd_helper.get_mesh_attrs(
+                prim, cache=self.usd_help._xform_cache, transform=robot_transform.numpy()
+            )
+            obstacles["mesh"].append(m)
+
+        for obj in self.robot.scene.objects:
+            if obj == self.robot:
+                continue
+            if obj.visual_only:
                 continue
-            prim_path_tokens = prim_path.split("/")
-            obj_name = prim_path_tokens[3]
-            link_name = prim_path_tokens[4]
-            mesh_name = prim_path_tokens[-1]
-            mesh = self.robot.scene.object_registry("name", obj_name).links[link_name].collision_meshes[mesh_name]
-            pos, orn = mesh.get_position_orientation()
-            inv_pos, inv_orn = T.invert_pose_transform(pos, orn)
-            # xyzw -> wxyz
-            inv_orn = inv_orn[[3, 0, 1, 2]]
-            inv_pose = self._tensor_args.to_device(th.cat([inv_pos, inv_orn]))
-            world_coll_checker._mesh_tensor_list[1][0, i, :7] = inv_pose
+            for link in obj.links.values():
+                for collision_mesh in link.collision_meshes.values():
+                    assert (
+                        collision_mesh.geom_type == "Mesh"
+                    ), f"collision_mesh {collision_mesh.prim_path} is not a mesh, but a {collision_mesh.geom_type}"
+                    obj_pose = T.pose2mat(collision_mesh.get_position_orientation())
+                    pose = robot_transform @ obj_pose
+                    pos, orn = T.mat2pose(pose)
+                    # xyzw -> wxyz
+                    orn = orn[[3, 0, 1, 2]]
+                    m = lazy.curobo.geom.types.Mesh(
+                        name=collision_mesh.prim_path,
+                        pose=th.cat([pos, orn]).tolist(),
+                        vertices=collision_mesh.points.numpy(),
+                        faces=collision_mesh.faces.numpy(),
+                        scale=collision_mesh.get_world_scale().numpy(),
+                    )
+                    obstacles["mesh"].append(m)
+
+        world = lazy.curobo.geom.types.WorldConfig(**obstacles)
+        world = world.get_collision_check_world()
+        self.mg[CuRoboEmbodimentSelection.DEFAULT].update_world(world)
 
     def check_collisions(
         self,

omnigibson/controllers/controller_base.py

@@ -364,12 +364,12 @@ def compute_no_op_action(self, control_dict):
         """
         if self._goal is None:
             self._goal = self.compute_no_op_goal(control_dict=control_dict)
-        command = self._compute_no_op_action(control_dict=control_dict)
+        command = self._compute_no_op_command(control_dict=control_dict)
         return cb.to_torch(self._reverse_preprocess_command(command))
 
-    def _compute_no_op_action(self, control_dict):
+    def _compute_no_op_command(self, control_dict):
         """
-        Compute no-op action given the goal
+        Compute no-op command given the goal
         """
         raise NotImplementedError
 

omnigibson/controllers/dd_controller.py

@@ -112,7 +112,7 @@ def compute_no_op_goal(self, control_dict):
         # This is zero-vector, since we want zero linear / angular velocity
         return dict(vel=cb.zeros(2))
 
-    def _compute_no_op_action(self, control_dict):
+    def _compute_no_op_command(self, control_dict):
         return cb.zeros(2)
 
     def _get_goal_shapes(self):

omnigibson/controllers/ik_controller.py

@@ -344,7 +344,7 @@ def compute_no_op_goal(self, control_dict):
             target_ori_mat=cb.as_float32(cb.T.quat2mat(control_dict[f"{self.task_name}_quat_relative"])),
         )
 
-    def _compute_no_op_action(self, control_dict):
+    def _compute_no_op_command(self, control_dict):
         pos_relative = control_dict[f"{self.task_name}_pos_relative"]
         quat_relative = control_dict[f"{self.task_name}_quat_relative"]
 

omnigibson/controllers/joint_controller.py

@@ -251,7 +251,7 @@ def compute_no_op_goal(self, control_dict):
 
         return dict(target=target)
 
-    def _compute_no_op_action(self, control_dict):
+    def _compute_no_op_command(self, control_dict):
         if self.motor_type == "position":
             if self._use_delta_commands:
                 return cb.zeros(self.command_dim)

omnigibson/controllers/multi_finger_gripper_controller.py

@@ -295,7 +295,7 @@ def compute_no_op_goal(self, control_dict):
         # Just use a zero vector
         return dict(target=cb.zeros(self.command_dim))
 
-    def _compute_no_op_action(self, control_dict):
+    def _compute_no_op_command(self, control_dict):
         # Take care of the special case of binary control
         if self._mode == "binary":
             command_val = -1 if self.is_grasping() == IsGraspingState.TRUE else 1

omnigibson/controllers/null_joint_controller.py

@@ -96,7 +96,7 @@ def update_default_goal(self, target):
 
         self.default_goal = cb.array(target)
 
-    def _compute_no_op_action(self, control_dict):
+    def _compute_no_op_command(self, control_dict):
         # Empty tensor since no action should be received
         return cb.array([])
 

omnigibson/controllers/osc_controller.py

@@ -457,7 +457,7 @@ def compute_no_op_goal(self, control_dict):
             target_ori_mat=cb.as_float32(cb.T.quat2mat(target_quat)),
         )
 
-    def _compute_no_op_action(self, control_dict):
+    def _compute_no_op_command(self, control_dict):
         pos_relative = control_dict[f"{self.task_name}_pos_relative"]
         quat_relative = control_dict[f"{self.task_name}_quat_relative"]
 

omnigibson/examples/robots/curobo_example.py

@@ -34,12 +34,12 @@ def execute_trajectory(q_traj, env, robot, attached_obj):
     for i, q in enumerate(q_traj):
         q = q.cpu()
         q = set_gripper_joint_positions(robot, q, attached_obj)
-        command = q_to_command(q, robot)
+        action = robot.q_to_action(q)
 
         num_repeat = 5
         print(f"Executing waypoint {i}/{len(q_traj)}")
         for _ in range(num_repeat):
-            env.step(command)
+            env.step(action)
 
 
 def plan_and_execute_trajectory(
@@ -76,21 +76,11 @@ def control_gripper(env, robot, attached_obj):
     # Control the gripper to open or close, while keeping the rest of the robot still
     q = robot.get_joint_positions()
     q = set_gripper_joint_positions(robot, q, attached_obj)
-    command = q_to_command(q, robot)
+    action = robot.q_to_action(q)
     num_repeat = 30
     print(f"Gripper (attached_obj={attached_obj})")
     for _ in range(num_repeat):
-        env.step(command)
-
-
-def q_to_command(q, robot):
-    # Convert target joint positions to command
-    command = []
-    for controller in robot.controllers.values():
-        command.append(q[controller.dof_idx])
-    command = th.cat(command, dim=0)
-    assert command.shape[0] == robot.action_dim
-    return command
+        env.step(action)
 
 
 def test_curobo():

omnigibson/objects/controllable_object.py

@@ -10,6 +10,7 @@
 import omnigibson as og
 from omnigibson.controllers import create_controller
 from omnigibson.controllers.controller_base import ControlType
+from omnigibson.controllers.joint_controller import JointController
 from omnigibson.objects.object_base import BaseObject
 from omnigibson.utils.backend_utils import _compute_backend as cb
 from omnigibson.utils.constants import JointType, PrimType
@@ -565,6 +566,12 @@ def deploy_control(self, control, control_type):
                 positions=control[pos_idxs],
                 indices=pos_idxs,
             )
+            # If we're setting joint position targets, we should also set velocity targets to 0
+            ControllableObjectViewAPI.set_joint_velocity_targets(
+                self.articulation_root_path,
+                velocities=cb.zeros(len(pos_idxs)),
+                indices=pos_idxs,
+            )
         vel_idxs = cb.where(control_type == ControlType.VELOCITY)[0]
         if len(vel_idxs) > 0:
             ControllableObjectViewAPI.set_joint_velocity_targets(
@@ -662,6 +669,24 @@ def _add_task_frame_control_dict(self, fcns, task_name, link_name):
             self.articulation_root_path
         )[-(self.n_links - link_idx), :, start_idx : start_idx + self.n_joints]
 
+    def q_to_action(self, q):
+        """
+        Converts a target joint configuration to an action that can be applied to this object.
+        All controllers should be JointController with use_delta_commands=False
+        """
+        action = []
+        for name, controller in self.controllers.items():
+            assert (
+                isinstance(controller, JointController) and not controller.use_delta_commands
+            ), f"Controller [{name}] should be a JointController with use_delta_commands=False!"
+            command = q[controller.dof_idx]
+            action.append(controller._reverse_preprocess_command(command))
+        action = th.cat(action, dim=0)
+        assert (
+            action.shape[0] == self.action_dim
+        ), f"Action should have dimension {self.action_dim}, got {action.shape[0]}"
+        return action
+
     def dump_action(self):
         """
         Dump the last action applied to this object. For use in demo collection.

omnigibson/prims/cloth_prim.py

@@ -136,6 +136,19 @@ def _post_load(self):
         # Store the default position of the points in the local frame
         self._default_positions = vtarray_to_torch(self.get_attribute(attr="points"))
 
+    # For cloth, points should NOT be @cached_property because their local poses change over time
+    @property
+    def points(self):
+        """
+        Returns:
+            th.tensor: Local poses of all points
+        """
+        # If the geom is a mesh we can directly return its points.
+        mesh = self.prim
+        mesh_type = mesh.GetPrimTypeInfo().GetTypeName()
+        assert mesh_type == "Mesh", f"Expected a mesh prim, got {mesh_type} instead!"
+        return vtarray_to_torch(mesh.GetAttribute("points").Get(), dtype=th.float32)
+
     @property
     def visual_aabb(self):
         return self.aabb

omnigibson/prims/entity_prim.py

@@ -1616,6 +1616,9 @@ def _dump_state(self):
             # We do NOT save joint pos / vel targets because this is only relevant for motorized joints (e.g.: robots).
             # Such control (a) only relies on the joint state, and not joint targets, when computing control, and
             # (b) these targets will be immediately overwritten as soon as the next physics step occurs.
+            # In other words, when loading the sim state, we will load the joint pos / vel and the controller state (e.g. goal)
+            # and then when we take a physics step, right before the physics step actually occurs, @og.sim._on_physics_step()
+            # will be called, which computes the control based on the just-loaded joint state and overwrites the joint targets with the new control.
             # So because these values are not used and require additional memory / compute, we do not save targets
 
         return state