Merge branch 'rl-primitives-training' of https://github.com/StanfordV…

…L/OmniGibson into rl-primitives-training
StanfordVL · Nov 17, 2023 · 4b9e86d · 4b9e86d
2 parents a74ee1d + 29c10b9
commit 4b9e86d
Show file tree

Hide file tree

Showing 25 changed files with 1,004 additions and 795 deletions.
diff --git a/omnigibson/action_primitives/starter_semantic_action_primitives.py b/omnigibson/action_primitives/starter_semantic_action_primitives.py
@@ -5,7 +5,7 @@
 See provided tiago_primitives.yaml config file for an example. See examples/action_primitives for
 runnable examples.
 """
-from functools import cached_property
+# from functools import cached_property
 import inspect
 import logging
 import random
@@ -828,7 +828,7 @@ def _target_in_reach_of_robot_relative(self, relative_target_pose):
         """
         return self._ik_solver_cartesian_to_joint_space(relative_target_pose) is not None
 
-    @cached_property
+    @property
     def _manipulation_control_idx(self):
         """The appropriate manipulation control idx for the current settings."""           
         if isinstance(self.robot, Tiago):
@@ -838,10 +838,13 @@ def _manipulation_control_idx(self):
             else:
                 return np.concatenate([self.robot.trunk_control_idx, self.robot.arm_control_idx[self.arm]])
 
+        if isinstance(self.robot, Fetch):
+            return np.concatenate([self.robot.trunk_control_idx, self.robot.arm_control_idx[self.arm]])
+
         # Otherwise just return the default arm control idx
         return self.robot.arm_control_idx[self.arm]
 
-    @cached_property
+    @property
     def _manipulation_descriptor_path(self):
         """The appropriate manipulation descriptor for the current settings."""           
         if isinstance(self.robot, Tiago) and m.TIAGO_TORSO_FIXED:

diff --git a/omnigibson/envs/rl_env.py b/omnigibson/envs/rl_env.py
diff --git a/omnigibson/robots/fetch.py b/omnigibson/robots/fetch.py
@@ -236,6 +236,11 @@ def _post_load(self):
             assert set(wheel_link.collision_meshes) == {"collisions"}, "Wheel link should only have 1 collision!"
             wheel_link.collision_meshes["collisions"].set_collision_approximation("boundingSphere")
 
+        # Temporarily enforce the base link to use a cube approximation (caster wheels are messed up otherwise)
+        base_link = self.root_link
+        assert set(base_link.collision_meshes) == {"collisions"}, "Base link should only have 1 collision!"
+        base_link.collision_meshes["collisions"].set_collision_approximation("boundingCube")
+
         # Also apply a convex decomposition to the torso lift link
         torso_lift_link = self.links["torso_lift_link"]
         assert set(torso_lift_link.collision_meshes) == {"collisions"}, "Wheel link should only have 1 collision!"

diff --git a/omnigibson/robots/robot_base.py b/omnigibson/robots/robot_base.py
@@ -269,7 +269,11 @@ def get_proprioception(self):
             n-array: numpy array of all robot-specific proprioceptive observations.
         """
         proprio_dict = self._get_proprioception_dict()
-        return np.concatenate([proprio_dict[obs] for obs in self._proprio_obs])
+        proprio = []
+        for obs_key in self._proprio_obs:
+            obs = proprio_dict[obs_key] if isinstance(proprio_dict[obs_key], np.ndarray) else np.array([proprio_dict[obs_key]]) 
+            proprio.append(obs)
+        return np.concatenate(proprio)
 
     def _get_proprioception_dict(self):
         """

diff --git a/omnigibson/tasks/grasp_task.py b/omnigibson/tasks/grasp_task.py
@@ -12,6 +12,7 @@
 from omnigibson.utils.motion_planning_utils import set_arm_and_detect_collision
 from omnigibson.utils.python_utils import classproperty
 from omnigibson.utils.sim_utils import land_object
+from omnigibson.action_primitives.starter_semantic_action_primitives import StarterSemanticActionPrimitives
 
 DIST_COEFF = 0.1
 GRASP_REWARD = 1.0
@@ -29,6 +30,7 @@ def __init__(
         reward_config=None,
     ):
         self.obj_name = obj_name
+        self._primitive_controller = None
         super().__init__(termination_config=termination_config, reward_config=reward_config)
 
     def _load(self, env):
@@ -53,65 +55,43 @@ def _create_reward_functions(self):
         return rewards
 
     def _reset_agent(self, env):
-        if hasattr(env, '_primitive_controller'):
-        #     robot = env.robots[0]
-        #     # Randomize the robots joint positions
-        #     joint_control_idx = np.concatenate([robot.trunk_control_idx, robot.arm_control_idx[robot.default_arm]])
-        #     joint_combined_idx = np.concatenate([robot.trunk_control_idx, robot.arm_control_idx["combined"]])
-        #     initial_joint_pos = np.array(robot.get_joint_positions()[joint_combined_idx])
-        #     control_idx_in_joint_pos = np.where(np.in1d(joint_combined_idx, joint_control_idx))[0]
-
-        #     with UndoableContext(env._primitive_controller.robot, env._primitive_controller.robot_copy, "original") as context:
-        #         for _ in range(MAX_JOINT_RANDOMIZATION_ATTEMPTS):
-        #             joint_pos, joint_control_idx = self._get_random_joint_position(robot)
-        #             initial_joint_pos[control_idx_in_joint_pos] = joint_pos
-        #             if not set_arm_and_detect_collision(context, initial_joint_pos):
-        #                 robot.set_joint_positions(joint_pos, joint_control_idx)
-        #                 og.sim.step()
-        #                 break
-
-            reset_pose_tiago = np.array([
-                -1.78029833e-04,  
-                3.20231302e-05, 
-                -1.85759447e-07,
-                0.0, 
-                -0.2,
-                0.0,  
-                0.1, 
-                -6.10000000e-01,
-                -1.10000000e+00,  
-                0.00000000e+00, 
-                -1.10000000e+00,  
-                1.47000000e+00,
-                0.00000000e+00,  
-                8.70000000e-01,  
-                2.71000000e+00,  
-                1.50000000e+00,
-                1.71000000e+00, 
-                -1.50000000e+00, 
-                -1.57000000e+00,  
-                4.50000000e-01,
-                1.39000000e+00,  
-                0.00000000e+00,  
-                0.00000000e+00,  
-                4.50000000e-02,
-                4.50000000e-02,  
-                4.50000000e-02,  
-                4.50000000e-02
-            ])
-            robot = env.robots[0]
-            robot.set_joint_positions(reset_pose_tiago)
-            og.sim.step()
-
-            # Randomize the robot's 2d pose
-            obj = env.scene.object_registry("name", self.obj_name)
-            grasp_poses = get_grasp_poses_for_object_sticky(obj)
-            grasp_pose, _ = random.choice(grasp_poses)
-            sampled_pose_2d = env._primitive_controller._sample_pose_near_object(obj, pose_on_obj=grasp_pose)
-            # sampled_pose_2d = [-0.433881, -0.210183, -2.96118]
-            robot_pose = env._primitive_controller._get_robot_pose_from_2d_pose(sampled_pose_2d)
-            robot.set_position_orientation(*robot_pose)
-            print("Reset robot pose to: ", robot_pose)
+        if self._primitive_controller is None:
+            self._primitive_controller = StarterSemanticActionPrimitives(env)
+
+        robot = env.robots[0]
+        # Randomize the robots joint positions
+        joint_control_idx = np.concatenate([robot.trunk_control_idx, robot.arm_control_idx[robot.default_arm]])
+        dim = len(joint_control_idx)
+        # For Tiago
+        if "combined" in robot.robot_arm_descriptor_yamls:
+            joint_combined_idx = np.concatenate([robot.trunk_control_idx, robot.arm_control_idx["combined"]])
+            initial_joint_pos = np.array(robot.get_joint_positions()[joint_combined_idx])
+            control_idx_in_joint_pos = np.where(np.in1d(joint_combined_idx, joint_control_idx))[0]
+        # For Fetch
+        else:
+            initial_joint_pos = np.array(robot.get_joint_positions()[joint_control_idx])
+            control_idx_in_joint_pos = np.arange(dim)
+
+        with PlanningContext(self._primitive_controller.robot, self._primitive_controller.robot_copy, "original") as context:
+            for _ in range(MAX_JOINT_RANDOMIZATION_ATTEMPTS):
+                joint_pos, joint_control_idx = self._get_random_joint_position(robot)
+                initial_joint_pos[control_idx_in_joint_pos] = joint_pos
+                if not set_arm_and_detect_collision(context, initial_joint_pos):
+                    robot.set_joint_positions(joint_pos, joint_control_idx)
+                    og.sim.step()
+                    break
+
+        # Randomize the robot's 2d pose
+        obj = env.scene.object_registry("name", self.obj_name)
+        grasp_poses = get_grasp_poses_for_object_sticky(obj)
+        grasp_pose, _ = random.choice(grasp_poses)
+        sampled_pose_2d = self._primitive_controller._sample_pose_near_object(obj, pose_on_obj=grasp_pose)
+        # sampled_pose_2d = [-0.433881, -0.210183, -2.96118]
+        robot_pose = self._primitive_controller._get_robot_pose_from_2d_pose(sampled_pose_2d)
+        robot.set_position_orientation(*robot_pose)
+
+        self._primitive_controller._settle_robot()
+        print("Reset robot pose to: ", robot_pose)
 
     # Overwrite reset by only removeing reset scene
     def reset(self, env):
@@ -122,7 +102,7 @@ def reset(self, env):
             env (Environment): environment instance to reset
         """
         # Reset the scene, agent, and variables
-        # self._reset_scene(env)
+        self._reset_scene(env)
         self._reset_agent(env)
         self._reset_variables(env)
 

diff --git a/rl/avg_metrics.py → rl/old_scripts/avg_metrics.py b/rl/avg_metrics.py → rl/old_scripts/avg_metrics.py
diff --git a/rl/cartpole_sample_training.py → rl/old_scripts/cartpole_sample_training.py b/rl/cartpole_sample_training.py → rl/old_scripts/cartpole_sample_training.py
diff --git a/rl/collect_data_rl.py → rl/old_scripts/collect_data_rl.py b/rl/collect_data_rl.py → rl/old_scripts/collect_data_rl.py
diff --git a/rl/training_bc.py → rl/old_scripts/training_bc.py b/rl/training_bc.py → rl/old_scripts/training_bc.py
diff --git a/rl/training_cql.py → rl/old_scripts/training_cql.py b/rl/training_cql.py → rl/old_scripts/training_cql.py
diff --git a/rl/training_sac.py → rl/old_scripts/training_sac.py b/rl/training_sac.py → rl/old_scripts/training_sac.py
diff --git a/rl/service/environment_pb2.py b/rl/service/environment_pb2.py