add AG info to simulator save & restore; resolve datatype mismatch wh…

…en restoring Tiago from json

omnigibson/robots/manipulation_robot.py

@@ -967,7 +967,7 @@ def _default_controller_config(self):
 
         return cfg
 
-    def _establish_grasp_rigid(self, arm="default", ag_data=None):
+    def _establish_grasp_rigid(self, arm="default", ag_data=None, contact_pos=None):
         """
         Establishes an ag-assisted grasp, if enabled.
 
@@ -976,6 +976,7 @@ def _establish_grasp_rigid(self, arm="default", ag_data=None):
                 Default is "default" which corresponds to the first entry in self.arm_names
             ag_data (None or 2-tuple): if specified, assisted-grasp object, link tuple (i.e. :(BaseObject, RigidPrim)).
                 Otherwise, does a no-op
+            contact_pos (None or np.array): if specified, contact position to use for grasp.
         """
         arm = self.default_arm if arm == "default" else arm
 
@@ -1001,12 +1002,12 @@ def _establish_grasp_rigid(self, arm="default", ag_data=None):
                 link_handle = self._dc.get_joint_parent_body(joint_handle)
                 joint_handle = self._dc.get_rigid_body_parent_joint(link_handle)
 
-        force_data, _ = self._find_gripper_contacts(arm=arm, return_contact_positions=True)
-        contact_pos = None
-        for c_link_prim_path, c_contact_pos in force_data:
-            if c_link_prim_path == ag_link.prim_path:
-                contact_pos = np.array(c_contact_pos)
-                break
+        if contact_pos is None:
+            force_data, _ = self._find_gripper_contacts(arm=arm, return_contact_positions=True)
+            for c_link_prim_path, c_contact_pos in force_data:
+                if c_link_prim_path == ag_link.prim_path:
+                    contact_pos = np.array(c_contact_pos)
+                    break
         assert contact_pos is not None
 
         # Joint frame set at the contact point
@@ -1048,6 +1049,7 @@ def _establish_grasp_rigid(self, arm="default", ag_data=None):
             "joint_type": joint_type,
             "gripper_pos": self.get_joint_positions()[self.gripper_control_idx[arm]],
             "max_force": max_force,
+            "contact_pos": contact_pos,
         }
         self._ag_obj_in_hand[arm] = ag_obj
         self._ag_freeze_gripper[arm] = True
@@ -1120,11 +1122,11 @@ def _calculate_in_hand_object(self, arm="default"):
         else:
             return self._calculate_in_hand_object_rigid(arm)
 
-    def _establish_grasp(self, arm="default", ag_data=None):
+    def _establish_grasp(self, arm="default", ag_data=None, contact_pos=None):
         if gm.AG_CLOTH:
             return self._establish_grasp_cloth(arm, ag_data)
         else:
-            return self._establish_grasp_rigid(arm, ag_data)
+            return self._establish_grasp_rigid(arm, ag_data, contact_pos)
 
     def _calculate_in_hand_object_cloth(self, arm="default"):
         """
@@ -1249,8 +1251,8 @@ def _dump_state(self):
         if self.grasping_mode == "physical":
             return state
 
-        # TODO: Include AG_state
-
+        # Include AG_state
+        state["ag_obj_constraint_params"] = self._ag_obj_constraint_params
         return state
 
     def _load_state(self, state):
@@ -1260,7 +1262,17 @@ def _load_state(self, state):
         if self.grasping_mode == "physical":
             return
 
-        # TODO: Include AG_state
+        # Include AG_state
+        # TODO: currently doese not take care of cloth objects
+        self._ag_obj_constraint_params = state["ag_obj_constraint_params"]
+        for arm in self._ag_obj_constraint_params.keys():
+            if len(self._ag_obj_constraint_params[arm]) > 0:
+                data = self._ag_obj_constraint_params[arm]
+                obj = og.sim.scene.object_registry("prim_path", data["ag_obj_prim_path"])
+                link = obj.links[data["ag_link_prim_path"].split("/")[-1]]
+                self._ag_data[arm] = (obj, link)
+                self._ag_obj_in_hand[arm] = obj
+                self._establish_grasp(arm=arm, ag_data=self._ag_data[arm], contact_pos=data["contact_pos"])
 
     def _serialize(self, state):
         # Call super first

omnigibson/robots/tiago.py

@@ -576,7 +576,13 @@ def set_position_orientation(self, position=None, orientation=None):
             position = current_position
         if orientation is None:
             orientation = current_orientation
-
+
+        # Position and orientation are lists when restoring scene from json. Cast to np.array
+        if isinstance(position, list):
+            position = np.array(position)
+        if isinstance(orientation, list):
+            orientation = np.array(orientation)
+
         # If the simulator is playing, set the 6 base joints to achieve the desired pose of base_footprint link frame
         if self._dc is not None and self._dc.is_simulating():
             # Find the relative transformation from base_footprint_link ("base_footprint") frame to root_link