Merge pull request #990 from StanfordVL/feat/controller-dependencies

Feat/controller dependencies

docs/tutorials/custom_robot_import.md

@@ -134,8 +134,8 @@ Now that we have the USD file for the robot, let's write our own robot class. Fo
             raise ValueError("Stretch does not support discrete actions!")
 
         @property
-        def controller_order(self):
-            # Controller ordering. Usually determined by general robot kinematics chain
+        def _raw_controller_order(self):
+            # Raw controller ordering. Usually determined by general robot kinematics chain
             # You can usually simply take a subset of these based on the type of robot interfaces inherited for your robot class
             return ["base", "camera", f"arm_{self.default_arm}", f"gripper_{self.default_arm}"]
 

omnigibson/configs/fetch_behavior.yaml

@@ -45,8 +45,6 @@ robots:
     action_normalize: true
     action_type: continuous
     grasping_mode: physical
-    rigid_trunk: false
-    default_trunk_offset: 0.365
     default_arm_pose: diagonal30
     default_reset_mode: tuck
     sensor_config:
@@ -61,8 +59,11 @@ robots:
     controller_config:
       base:
         name: DifferentialDriveController
+      trunk:
+        name: JointController
       arm_0:
         name: InverseKinematicsController
+        subsume_controllers: [trunk]
       gripper_0:
         name: MultiFingerGripperController
         mode: binary

omnigibson/configs/fetch_primitives.yaml

@@ -38,14 +38,15 @@ robots:
     action_normalize: false
     action_type: continuous
     grasping_mode: sticky
-    rigid_trunk: false
-    default_trunk_offset: 0.365
     default_arm_pose: diagonal30
     controller_config:
       base:
         name: DifferentialDriveController
+      trunk:
+        name: JointController
       arm_0:
         name: JointController
+        subsume_controllers: [trunk]
         motor_type: position
         command_input_limits: null
         use_delta_commands: false

omnigibson/configs/robots/fetch.yaml

@@ -16,8 +16,6 @@ robot:
   scale: 1.0
   self_collision: true
   grasping_mode: physical
-  rigid_trunk: false
-  default_trunk_offset: 0.365
   default_arm_pose: vertical
   sensor_config:
     VisionSensor:
@@ -31,8 +29,11 @@ robot:
   controller_config:
     base:
       name: DifferentialDriveController
+    trunk:
+      name: JointController
     arm_0:
       name: InverseKinematicsController
+      subsume_controllers: [trunk]
     gripper_0:
       name: MultiFingerGripperController
     camera:

omnigibson/configs/tiago_primitives.yaml

@@ -38,8 +38,6 @@ robots:
     action_normalize: false
     action_type: continuous
     grasping_mode: sticky
-    rigid_trunk: false
-    default_trunk_offset: 0.15
     default_arm_pose: vertical
     sensor_config:
       VisionSensor:
@@ -49,6 +47,8 @@ robots:
     controller_config:
       base:
         name: JointController
+      trunk:
+        name: JointController
       arm_left:
         name: JointController
         motor_type: position
@@ -57,6 +57,7 @@ robots:
         use_delta_commands: false
       arm_right:
         name: JointController
+        subsume_controllers: [trunk]
         motor_type: position
         command_input_limits: null
         command_output_limits: null

omnigibson/controllers/ik_controller.py

@@ -38,9 +38,6 @@ class InverseKinematicsController(JointController, ManipulationController):
     def __init__(
         self,
         task_name,
-        robot_description_path,
-        robot_urdf_path,
-        eef_name,
         control_freq,
         reset_joint_pos,
         control_limits,
@@ -63,9 +60,6 @@ def __init__(
             task_name (str): name assigned to this task frame for computing IK control. During control calculations,
                 the inputted control_dict should include entries named <@task_name>_pos_relative and
                 <@task_name>_quat_relative. See self._command_to_control() for what these values should entail.
-            robot_description_path (str): path to robot descriptor yaml file
-            robot_urdf_path (str): path to robot urdf file
-            eef_name (str): end effector frame name
             control_freq (int): controller loop frequency
             reset_joint_pos (Array[float]): reset joint positions, used as part of nullspace controller in IK.
                 Note that this should correspond to ALL the joints; the exact indices will be extracted via @dof_idx

omnigibson/objects/controllable_object.py

@@ -231,10 +231,44 @@ def _load_controllers(self):
 
         # Initialize controllers to create
         self._controllers = dict()
-        # Loop over all controllers, in the order corresponding to @action dim
-        for name in self.controller_order:
+        # Keep track of any controllers that are subsumed by other controllers
+        # We will not instantiate subsumed controllers
+        controller_subsumes = dict()  # Maps independent controller name to list of subsumed controllers
+        subsume_names = set()
+        for name in self._raw_controller_order:
+            # Make sure we have the valid controller name specified
             assert_valid_key(key=name, valid_keys=self._controller_config, name="controller name")
             cfg = self._controller_config[name]
+            subsume_controllers = cfg.pop("subsume_controllers", [])
+            # If this controller subsumes other controllers, it cannot be subsumed by another controller
+            # (i.e.: we don't allow nested / cyclical subsuming)
+            if len(subsume_controllers) > 0:
+                assert (
+                    name not in subsume_names
+                ), f"Controller {name} subsumes other controllers, and therefore cannot be subsumed by another controller!"
+                controller_subsumes[name] = subsume_controllers
+                for subsume_name in subsume_controllers:
+                    # Make sure it doesn't already exist -- a controller should only be subsumed by up to one other
+                    assert (
+                        subsume_name not in subsume_names
+                    ), f"Controller {subsume_name} cannot be subsumed by more than one other controller!"
+                    assert (
+                        subsume_name not in controller_subsumes
+                    ), f"Controller {name} subsumes other controllers, and therefore cannot be subsumed by another controller!"
+                    subsume_names.add(subsume_name)
+
+        # Loop over all controllers, in the order corresponding to @action dim
+        for name in self._raw_controller_order:
+            # If this controller is subsumed by another controller, simply skip it
+            if name in subsume_names:
+                continue
+            cfg = self._controller_config[name]
+            # If we subsume other controllers, prepend the subsumed' dof idxs to this controller's idxs
+            if name in controller_subsumes:
+                for subsumed_name in controller_subsumes[name]:
+                    subsumed_cfg = self._controller_config[subsumed_name]
+                    cfg["dof_idx"] = th.concatenate([subsumed_cfg["dof_idx"], cfg["dof_idx"]])
+
             # If we're using normalized action space, override the inputs for all controllers
             if self._action_normalize:
                 cfg["command_input_limits"] = "default"  # default is normalized (-1, 1)
@@ -254,8 +288,12 @@ def update_controller_mode(self):
         Helper function to force the joints to use the internal specified control mode and gains
         """
         # Update the control modes of each joint based on the outputted control from the controllers
+        unused_dofs = {i for i in range(self.n_dof)}
         for name in self._controllers:
             for dof in self._controllers[name].dof_idx:
+                # Make sure the DOF has not already been set yet, and remove it afterwards
+                assert dof.item() in unused_dofs
+                unused_dofs.remove(dof.item())
                 control_type = self._controllers[name].control_type
                 self._joints[self.dof_names_ordered[dof]].set_control_type(
                     control_type=control_type,
@@ -267,6 +305,21 @@ def update_controller_mode(self):
                     ),
                 )
 
+        # For all remaining DOFs not controlled, we assume these are free DOFs (e.g.: virtual joints representing free
+        # motion wrt a specific axis), so explicitly set kp / kd to 0 to avoid silent bugs when
+        # joint positions / velocities are set
+        for unused_dof in unused_dofs:
+            unused_joint = self._joints[self.dof_names_ordered[unused_dof]]
+            assert not unused_joint.driven, (
+                f"All unused joints not mapped to any controller should not have DriveAPI attached to it! "
+                f"However, joint {unused_joint.name} is driven!"
+            )
+            unused_joint.set_control_type(
+                control_type=ControlType.NONE,
+                kp=None,
+                kd=None,
+            )
+
     def _generate_controller_config(self, custom_config=None):
         """
         Generates a fully-populated controller config, overriding any default values with the corresponding values
@@ -283,7 +336,7 @@ def _generate_controller_config(self, custom_config=None):
         controller_config = {} if custom_config is None else deepcopy(custom_config)
 
         # Update the configs
-        for group in self.controller_order:
+        for group in self._raw_controller_order:
             group_controller_name = (
                 controller_config[group]["name"]
                 if group in controller_config and "name" in controller_config[group]
@@ -623,6 +676,41 @@ def get_control_dict(self):
 
         return fcns
 
+    def _add_task_frame_control_dict(self, fcns, task_name, link_name):
+        """
+        Internally helper function to generate per-link control dictionary entries. Useful for generating relevant
+        control values needed for IK / OSC for a given @task_name. Should be called within @get_control_dict()
+
+        Args:
+            fcns (CachedFunctions): Keyword-mapped control values for this object, mapping names to n-arrays.
+            task_name (str): name to assign for this task_frame. It will be prepended to all fcns generated
+            link_name (str): the corresponding link name from this controllable object that @task_name is referencing
+        """
+        fcns[f"_{task_name}_pos_quat_relative"] = (
+            lambda: ControllableObjectViewAPI.get_link_relative_position_orientation(
+                self.articulation_root_path, link_name
+            )
+        )
+        fcns[f"{task_name}_pos_relative"] = lambda: fcns[f"_{task_name}_pos_quat_relative"][0]
+        fcns[f"{task_name}_quat_relative"] = lambda: fcns[f"_{task_name}_pos_quat_relative"][1]
+        fcns[f"{task_name}_lin_vel_relative"] = lambda: ControllableObjectViewAPI.get_link_relative_linear_velocity(
+            self.articulation_root_path, link_name
+        )
+        fcns[f"{task_name}_ang_vel_relative"] = lambda: ControllableObjectViewAPI.get_link_relative_angular_velocity(
+            self.articulation_root_path, link_name
+        )
+        # -n_joints because there may be an additional 6 entries at the beginning of the array, if this robot does
+        # not have a fixed base (i.e.: the 6DOF --> "floating" joint)
+        # see self.get_relative_jacobian() for more info
+        # We also count backwards for the link frame because if the robot is fixed base, the jacobian returned has one
+        # less index than the number of links. This is presumably because the 1st link of a fixed base robot will
+        # always have a zero jacobian since it can't move. Counting backwards resolves this issue.
+        start_idx = 0 if self.fixed_base else 6
+        link_idx = self._articulation_view.get_body_index(link_name)
+        fcns[f"{task_name}_jacobian_relative"] = lambda: ControllableObjectViewAPI.get_relative_jacobian(
+            self.articulation_root_path
+        )[-(self.n_links - link_idx), :, start_idx : start_idx + self.n_joints]
+
     def dump_action(self):
         """
         Dump the last action applied to this object. For use in demo collection.
@@ -755,13 +843,27 @@ def controllers(self):
         return self._controllers
 
     @property
-    @abstractmethod
     def controller_order(self):
         """
         Returns:
             list: Ordering of the actions, corresponding to the controllers. e.g., ["base", "arm", "gripper"],
                 to denote that the action vector should be interpreted as first the base action, then arm command, then
-                gripper command
+                gripper command. Note that this may be a subset of all possible controllers due to some controllers
+                subsuming others (e.g.: arm controller subsuming the trunk controller if using IK)
+        """
+        assert self._controllers is not None, "Can only view controller_order after controllers are loaded!"
+        return list(self._controllers.keys())
+
+    @property
+    @abstractmethod
+    def _raw_controller_order(self):
+        """
+        Returns:
+            list: Raw ordering of the actions, corresponding to the controllers. e.g., ["base", "arm", "gripper"],
+                to denote that the action vector should be interpreted as first the base action, then arm command, then
+                gripper command. Note that external users should query @controller_order, which is the post-processed
+                ordering of actions, which may be a subset of the controllers due to some controllers subsuming others
+                (e.g.: arm controller subsuming the trunk controller if using IK)
         """
         raise NotImplementedError
 

omnigibson/prims/entity_prim.py

@@ -1612,7 +1612,6 @@ def _dump_state(self):
         if self.n_joints > 0:
             state["joint_pos"] = self.get_joint_positions()
             state["joint_vel"] = self.get_joint_velocities()
-            state["joint_eff"] = self.get_joint_efforts()
 
             # 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
@@ -1633,7 +1632,6 @@ def _load_state(self, state):
         elif self.n_joints > 0:
             self.set_joint_positions(state["joint_pos"])
             self.set_joint_velocities(state["joint_vel"])
-            self.set_joint_efforts(state["joint_eff"])
 
         # Make sure this object is awake
         self.wake()
@@ -1646,7 +1644,6 @@ def serialize(self, state):
             state_flat += [
                 state["joint_pos"],
                 state["joint_vel"],
-                state["joint_eff"],
             ]
 
         return th.cat(state_flat)
@@ -1657,7 +1654,7 @@ def deserialize(self, state):
         root_link_state, idx = self.root_link.deserialize(state=state)
         state_dict = dict(root_link=root_link_state)
         if self.n_joints > 0:
-            for jnt_state in ("pos", "vel", "eff"):
+            for jnt_state in ("pos", "vel"):
                 state_dict[f"joint_{jnt_state}"] = state[idx : idx + self.n_joints]
                 idx += self.n_joints
 

omnigibson/prims/geom_prim.py

@@ -117,7 +117,7 @@ def opacity(self, opacity):
         if self.has_material():
             self.material.opacity_constant = opacity
         else:
-            self.set_attribute("primvars:displayOpacity", th.tensor([opacity]))
+            self.set_attribute("primvars:displayOpacity", np.array([opacity]))
 
     @property
     def points(self):

omnigibson/prims/joint_prim.py

@@ -173,11 +173,12 @@ def update_handles(self):
 
     def set_control_type(self, control_type, kp=None, kd=None):
         """
-        Sets the control type for this joint.
+        Sets the control type for this joint. Note that ControlType.NONE is equivalent to
+        ControlType.EFFORT with 0 kp / kd
 
         Args:
             control_type (ControlType): What type of control to use for this joint.
-                Valid options are: {ControlType.POSITION, ControlType.VELOCITY, ControlType.EFFORT}
+                Valid options are: {ControlType.POSITION, ControlType.VELOCITY, ControlType.EFFORT, ControlType.NONE}
             kp (None or float): If specified, sets the kp gain value for this joint. Should only be set if
                 setting ControlType.POSITION
             kd (None or float): If specified, sets the kd gain value for this joint. Should only be set if
@@ -194,7 +195,7 @@ def set_control_type(self, control_type, kp=None, kd=None):
             assert kp is None, "kp gain must not be specified for setting VELOCITY control!"
             assert kd is not None, "kd gain must be specified for setting VELOCITY control!"
             kp = 0.0
-        else:  # Efforts
+        else:  # Efforts (or NONE -- equivalent)
             assert kp is None, "kp gain must not be specified for setting EFFORT control!"
             assert kd is None, "kd gain must not be specified for setting EFFORT control!"
             kp, kd = 0.0, 0.0
@@ -867,12 +868,11 @@ def keep_still(self):
             self.set_effort(th.zeros(self.n_dof))
 
     def _dump_state(self):
-        pos, vel, effort = self.get_state() if self.articulated else (th.empty(0), th.empty(0), th.empty(0))
+        pos, vel, _ = self.get_state() if self.articulated else (th.empty(0), th.empty(0), th.empty(0))
         target_pos, target_vel = self.get_target() if self.articulated else (th.empty(0), th.empty(0))
         return dict(
             pos=pos,
             vel=vel,
-            effort=effort,
             target_pos=target_pos,
             target_vel=target_vel,
         )
@@ -881,8 +881,6 @@ def _load_state(self, state):
         if self.articulated:
             self.set_pos(state["pos"], drive=False)
             self.set_vel(state["vel"], drive=False)
-            if self.driven:
-                self.set_effort(state["effort"])
             if self._control_type == ControlType.POSITION:
                 self.set_pos(state["target_pos"], drive=True)
             elif self._control_type == ControlType.VELOCITY:
@@ -893,7 +891,6 @@ def serialize(self, state):
             [
                 state["pos"],
                 state["vel"],
-                state["effort"],
                 state["target_pos"],
                 state["target_vel"],
             ]
@@ -905,9 +902,8 @@ def deserialize(self, state):
             dict(
                 pos=state[0 : self.n_dof],
                 vel=state[self.n_dof : 2 * self.n_dof],
-                effort=state[2 * self.n_dof : 3 * self.n_dof],
-                target_pos=state[3 * self.n_dof : 4 * self.n_dof],
-                target_vel=state[4 * self.n_dof : 5 * self.n_dof],
+                target_pos=state[2 * self.n_dof : 3 * self.n_dof],
+                target_vel=state[3 * self.n_dof : 4 * self.n_dof],
             ),
-            5 * self.n_dof,
+            4 * self.n_dof,
         )