Integrate cuRobo into action primitives & primitives refactoring

.github/workflows/tests.yml

@@ -60,7 +60,7 @@ jobs:
         continue-on-error: true
 
       - name: Deploy artifact
-        uses: actions/upload-artifact@v3
+        uses: actions/upload-artifact@v4
         with:
           name: ${{ github.run_id }}-tests-${{ matrix.test_file }}
           path: ${{ matrix.test_file }}.xml

docker/prod.Dockerfile

@@ -44,7 +44,7 @@ RUN wget --no-verbose -O /cuda-keyring.deb https://developer.download.nvidia.com
   DEBIAN_FRONTEND=noninteractive apt-get install -y cuda-toolkit-11-8 && \
   TORCH_CUDA_ARCH_LIST='7.5;8.0;8.6+PTX' PATH=/usr/local/cuda-11.8/bin:$PATH LD_LIBRARY_PATH=/usr/local/cuda-11.8/lib64:$LD_LIBRARY_PATH \
     micromamba run -n omnigibson pip install \
-    git+https://github.com/StanfordVL/curobo@6a4eb2ca8677829b0f57451ad107e0a3186525e9#egg=nvidia_curobo \
+    git+https://github.com/StanfordVL/curobo@cbaf7d32436160956dad190a9465360fad6aba73#egg=nvidia_curobo \
     --no-build-isolation > /dev/null && \
   apt-get remove -y cuda-toolkit-11-8 && apt-get autoremove -y && apt-get autoclean -y && rm -rf /var/lib/apt/lists/*
 

omnigibson/action_primitives/action_primitive_set_base.py

@@ -58,14 +58,9 @@ def __init_subclass__(cls, **kwargs):
         if not inspect.isabstract(cls):
             REGISTERED_PRIMITIVE_SETS[cls.__name__] = cls
 
-    def __init__(self, env):
+    def __init__(self, env, robot):
         self.env = env
-
-    @property
-    def robot(self):
-        # Currently returns the first robot in the environment, but can be scaled to multiple robots
-        # by creating multiple action generators and passing in a robot index etc.
-        return self.env.robots[0]
+        self.robot = robot
 
     @abstractmethod
     def get_action_space(self):
@@ -84,7 +79,7 @@ def apply(self, action):
         pass
 
     @abstractmethod
-    def apply_ref(self, action, *args):
+    def apply_ref(self, primitive, *args):
         """
         Apply a primitive action by reference.
 

omnigibson/action_primitives/symbolic_semantic_action_primitives.py

@@ -40,8 +40,8 @@ class SymbolicSemanticActionPrimitiveSet(IntEnum):
 
 
 class SymbolicSemanticActionPrimitives(StarterSemanticActionPrimitives):
-    def __init__(self, env):
-        super().__init__(env)
+    def __init__(self, env, robot):
+        super().__init__(env, robot, skip_curobo_initilization=True)
         self.controller_functions = {
             SymbolicSemanticActionPrimitiveSet.GRASP: self._grasp,
             SymbolicSemanticActionPrimitiveSet.PLACE_ON_TOP: self._place_on_top,
@@ -59,12 +59,12 @@ def __init__(self, env):
             SymbolicSemanticActionPrimitiveSet.RELEASE: self._release,
         }
 
-    def apply_ref(self, prim, *args, attempts=3):
+    def apply_ref(self, primitive, *args, attempts=3):
         """
         Yields action for robot to execute the primitive with the given arguments.
 
         Args:
-            prim (SymbolicSemanticActionPrimitiveSet): Primitive to execute
+            primitive (SymbolicSemanticActionPrimitiveSet): Primitive to execute
             args: Arguments for the primitive
             attempts (int): Number of attempts to make before raising an error
 
@@ -75,7 +75,7 @@ def apply_ref(self, prim, *args, attempts=3):
             ActionPrimitiveError: If primitive fails to execute
         """
         assert attempts > 0, "Must make at least one attempt"
-        ctrl = self.controller_functions[prim]
+        ctrl = self.controller_functions[primitive]
 
         if any(isinstance(arg, BaseRobot) for arg in args):
             raise ActionPrimitiveErrorGroup(

omnigibson/configs/fetch_primitives.yaml → omnigibson/configs/r1_primitives.yaml

@@ -6,7 +6,6 @@ env:
   flatten_action_space: false           # (bool): whether to flatten the action space as a sinle 1D-array
   flatten_obs_space: false              # (bool): whether the observation space should be flattened when generated
   use_external_obs: false               # (bool): Whether to use external observations or not
-  initial_pos_z_offset: 0.1
   external_sensors: null                # (None or list): If specified, list of sensor configurations for external sensors to add. Should specify sensor "type" and any additional kwargs to instantiate the sensor. Each entry should be the kwargs passed to @create_sensor, in addition to position, orientation
 
 render:
@@ -31,36 +30,83 @@ scene:
   include_robots: false
 
 robots:
-  - type: Fetch
+  - type: R1
     obs_modalities: [rgb]
     scale: 1.0
     self_collisions: true
     action_normalize: false
     action_type: continuous
     grasping_mode: sticky
-    default_arm_pose: diagonal30
+    reset_joint_pos: [
+      0.0000,
+      -0.0000,
+      0.0247,
+      0.0009,
+      0.0004,
+      -0.0000,
+      0.0000,
+      0.0000,
+      0.0000,
+      0.0000,
+      -0.0464,
+      0.0464,
+      2.6172,
+      2.6168,
+      -1.4584,
+      -1.4570,
+      -0.0433,
+      0.0418,
+      1.5899,
+      -1.5896,
+      -1.1587,
+      1.1593,
+      0.0300,
+      0.0300,
+      0.0300,
+      0.0300,
+    ]
+    sensor_config:
+      VisionSensor:
+        sensor_kwargs:
+          image_height: 128
+          image_width: 128
     controller_config:
       base:
-        name: DifferentialDriveController
+        name: HolonomicBaseJointController
+        motor_type: position
+        command_input_limits: null
+        use_impedances: false
       trunk:
         name: JointController
-      arm_0:
+        motor_type: position
+        command_input_limits: null
+        use_delta_commands: false
+        use_impedances: false
+      arm_left:
+        name: JointController
+        motor_type: position
+        command_input_limits: null
+        use_delta_commands: false
+        use_impedances: false
+      arm_right:
         name: JointController
-        subsume_controllers: [trunk]
         motor_type: position
         command_input_limits: null
         use_delta_commands: false
-      gripper_0:
+        use_impedances: false
+      gripper_left:
         name: JointController
         motor_type: position
-        command_input_limits: [-1, 1]
-        command_output_limits: null
+        command_input_limits: null
         use_delta_commands: false
-      camera:
+        use_impedances: false
+      gripper_right:
         name: JointController
-        use_delta_commands: False
+        motor_type: position
+        command_input_limits: null
+        use_delta_commands: false
+        use_impedances: false
 
 objects: []
-
 task:
   type: DummyTask

omnigibson/configs/tiago_primitives.yaml

@@ -32,51 +32,88 @@ scene:
 
 robots:
   - type: Tiago
-    obs_modalities: [rgb, depth, seg_semantic, normal, seg_instance, seg_instance_id]
+    obs_modalities: [rgb]
     scale: 1.0
     self_collisions: true
     action_normalize: false
     action_type: continuous
     grasping_mode: sticky
-    default_arm_pose: vertical
+    reset_joint_pos: [
+      0.0000,
+      0.0000,
+      -0.0000,
+      -0.0000,
+      -0.0000,
+      -0.0000,
+      0.3500,
+      0.9052,
+      0.9052,
+      0.0000,
+      -0.4281,
+      -0.4281,
+      -0.4500,
+      2.2350,
+      2.2350,
+      1.6463,
+      1.6463,
+      0.7687,
+      0.7687,
+      -0.7946,
+      -0.7946,
+      -1.0891,
+      -1.0891,
+      0.0450,
+      0.0450,
+      0.0450,
+      0.0450,
+    ]
     sensor_config:
       VisionSensor:
         sensor_kwargs:
           image_height: 128
           image_width: 128
     controller_config:
       base:
-        name: JointController
+        name: HolonomicBaseJointController
+        motor_type: position
+        command_input_limits: null
+        use_impedances: false
       trunk:
         name: JointController
+        motor_type: position
+        command_input_limits: null
+        use_delta_commands: false
+        use_impedances: false
       arm_left:
         name: JointController
-        subsume_controllers: [trunk]
         motor_type: position
         command_input_limits: null
-        command_output_limits: null
         use_delta_commands: false
+        use_impedances: False
       arm_right:
         name: JointController
         motor_type: position
         command_input_limits: null
-        command_output_limits: null
         use_delta_commands: false
+        use_impedances: false
       gripper_left:
         name: JointController
         motor_type: position
-        command_input_limits: [-1, 1]
-        command_output_limits: null
+        command_input_limits: null
         use_delta_commands: false
+        use_impedances: false
       gripper_right:
         name: JointController
         motor_type: position
-        command_input_limits: [-1, 1]
-        command_output_limits: null
+        command_input_limits: null
         use_delta_commands: false
+        use_impedances: false
       camera:
         name: JointController
-
+        motor_type: position
+        command_input_limits: null
+        use_delta_commands: false
+        use_impedances: false
 objects: []
 
 task:

omnigibson/controllers/__init__.py

@@ -9,6 +9,7 @@
     ManipulationController,
 )
 from omnigibson.controllers.dd_controller import DifferentialDriveController
+from omnigibson.controllers.holonomic_base_joint_controller import HolonomicBaseJointController
 from omnigibson.controllers.ik_controller import InverseKinematicsController
 from omnigibson.controllers.joint_controller import JointController
 from omnigibson.controllers.multi_finger_gripper_controller import MultiFingerGripperController
@@ -39,6 +40,7 @@ def create_controller(name, **kwargs):
     "create_controller",
     "DifferentialDriveController",
     "GripperController",
+    "HolonomicBaseJointController",
     "InverseKinematicsController",
     "IsGraspingState",
     "JointController",