add trunk action idx

StanfordVL · Oct 18, 2024 · 15c58c0 · 15c58c0
1 parent 94deb2e
commit 15c58c0
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/omnigibson/robots/articulated_trunk_robot.py b/omnigibson/robots/articulated_trunk_robot.py
@@ -45,6 +45,12 @@ def trunk_control_idx(self):
         """
         return th.tensor([list(self.joints.keys()).index(name) for name in self.trunk_joint_names])
 
+    @property
+    def trunk_action_idx(self):
+        controller_idx = self.controller_order.index("trunk")
+        action_start_idx = sum([self.controllers[self.controller_order[i]].command_dim for i in range(controller_idx)])
+        return th.arange(action_start_idx, action_start_idx + self.controllers["trunk"].command_dim)
+
     @property
     def _default_controllers(self):
         # Always call super first