Merge pull request #3 from StanfordVL/gripper_and_objects

Gripper and objects

MujocoManip/environment/base.py

@@ -66,8 +66,8 @@ def _get_reference(self):
         pass
 
     def reset(self):
-        with self.physics.reset_context():
-            self._reset_internal()
+        self._reset_internal()
+        self.physics.forward()
         return self._get_observation()
 
     def _reset_internal(self):

MujocoManip/environment/sawyer.py

@@ -20,8 +20,8 @@ def __init__(self, gripper=None, use_eef_ctrl=False, use_torque_ctrl=False, use_
             self._setup_mocap()
 
     def _setup_mocap(self):
-        with self.physics.reset_context():
-            reset_mocap_welds(self.physics)
+        reset_mocap_welds(self.physics)
+        self.physics.forward()
 
         # Move end effector into position.
         gripper_target = self.physics.named.data.xpos['right_hand']
@@ -166,8 +166,8 @@ def set_robot_joint_positions(self, jpos):
         """
         Helper method to force robot joint positions to the passed values.
         """
-        with self.physics.reset_context():
-            self.physics.named.data.qpos[self.mujoco_robot.joints] = jpos
+        self.physics.named.data.qpos[self.mujoco_robot.joints] = jpos
+        self.physics.forward()
 
     @property
     #TODO: fix it

MujocoManip/environment/sawyer_stack.py

@@ -9,7 +9,7 @@ class SawyerStackEnv(SawyerEnv):
     def __init__(self, 
                 gripper='TwoFingerGripper',
                 mujoco_objects=None,
-                n_mujoco_objects=5,
+                n_mujoco_objects=10,
                 table_size=(0.8, 0.8, 0.8),
                 table_friction=None,
                 reward_lose=-2,
@@ -35,7 +35,7 @@ def __init__(self,
         # Handle parameters
         self.mujoco_objects = mujoco_objects
         if self.mujoco_objects is None:
-            self.mujoco_objects = [RandomBoxObject() for i in range(n_mujoco_objects)]
+            self.mujoco_objects = [RandomBoxObject(size_max=[0.03, 0.03, 0.03], size_min=[0.01, 0.01, 0.01]) for i in range(n_mujoco_objects)]
         self.n_mujoco_objects = len(self.mujoco_objects)
         self.table_size = table_size
         self.table_friction = table_friction
@@ -93,32 +93,44 @@ def _place_targets(self):
             contact_point = contact_point - di['object_bottom_offset'] + di['object_top_offset']
 
     def _place_objects(self):
+        """
+        Place objects randomly until no more collisions or max iterations hit.
+        """
         placed_objects = []
         for index in range(self.n_objects):
             di = self.object_metadata[index]
             horizontal_radius = di['object_horizontal_radius']
+            bottom_offset = di['object_bottom_offset']
             success = False
-            for i in range(1000): # 1000 retries
+            for i in range(100000): # 1000 retries
                 object_x = np.random.uniform(high=self.table_size[0]/2 - horizontal_radius, low=-1 * (self.table_size[0]/2 - horizontal_radius))
                 object_y = np.random.uniform(high=self.table_size[1]/2 - horizontal_radius, low=-1 * (self.table_size[1]/2 - horizontal_radius))
                 # objects cannot overlap
                 location_valid = True
-                for x, y, r in placed_objects:
-                    if np.linalg.norm([object_x - x, object_y - y], 2) < r + horizontal_radius:
+                for _, (x, y, z), r in placed_objects:
+                    if np.linalg.norm([object_x - x, object_y - y], 2) <= r + horizontal_radius:
                         location_valid = False
                         break
                 if not location_valid: # bad luck, reroll
                     continue
                 # location is valid, put the object down
                 # quarternions, later we can add random rotation
+
                 pos = np.array([object_x, object_y, 0])
-                self._set_object_pos(index, pos)
-                self._set_object_vel(index, np.zeros(3))
-                placed_objects.append((object_x, object_y, horizontal_radius))
+                pos -= bottom_offset
+                placed_objects.append((index, pos, horizontal_radius))
                 success = True
                 break
             if not success:
                 raise RandomizationError('Cannot place all objects on the desk')
+
+        # with self.physics.reset_context():
+        for index, pos, r in placed_objects:
+            self._set_object_pos(index, pos)
+            self._set_object_vel(index, np.zeros(3))
+        self.physics.forward()
+            # pos = np.array([object_x, object_y, 0])
+
 
     def _reward(self, action):
         reward = 0
@@ -175,8 +187,7 @@ def _object_pos(self, i):
         return self.physics.named.data.xpos[object_name] - self._pos_offset
 
     def _set_object_pos(self, i, pos):
-        pos += self._pos_offset
-        self.physics.named.data.qpos[self.object_metadata[i]['joint_name']][0:3] = pos 
+        self.physics.named.data.qpos[self.object_metadata[i]['joint_name']][0:3] = pos + self._pos_offset
 
     def _object_vel(self, i):
         object_name = self.object_metadata[i]['object_name']

MujocoManip/miscellaneous/utils.py

@@ -466,12 +466,12 @@ def mocap_set_action(physics, action):
 def reset_mocap_welds(physics):
     """Resets the mocap welds that we use for actuation.
     """
-    with physics.reset_context():
-        if physics.model.nmocap > 0 and physics.model.eq_data is not None:
-            for i in range(physics.model.eq_data.shape[0]):
-                if physics.model.eq_type[i] == enums.mjtEq.mjEQ_WELD:
-                    physics.model.eq_data[i, :] = np.array(
-                        [0., 0., 0., 1., 0., 0., 0.])
+    if physics.model.nmocap > 0 and physics.model.eq_data is not None:
+        for i in range(physics.model.eq_data.shape[0]):
+            if physics.model.eq_type[i] == enums.mjtEq.mjEQ_WELD:
+                physics.model.eq_data[i, :] = np.array(
+                    [0., 0., 0., 1., 0., 0., 0.])
+    physics.forward()
 
 
 def reset_mocap2body_xpos(physics):

MujocoManip/model/assets/base.xml

@@ -1,6 +1,7 @@
+<!-- This is the base xml for all physics simulations. Set global configs here. -->
 <mujoco model="base">
 	<compiler angle="radian" meshdir="meshes/"/>
-	<size nconmax="200" njmax="1000"/>
+	<size nconmax="1000" njmax="1000"/>
 
 	<asset>
 	</asset>

MujocoManip/model/assets/gripper/two_finger_gripper.xml

@@ -6,9 +6,7 @@
     </asset>
 
     <actuator>
-        <!-- TODO: play with range -->
         <position ctrllimited="true" ctrlrange="-0.008 0.020833" joint="r_gripper_r_finger_joint" kp="10000" name="gripper_r_gripper_r_finger_joint" forcelimited="true" forcerange="-20 20"/>
-        <!-- TODO: play with range -->
         <position ctrllimited="true" ctrlrange="-0.020833 0.008" joint="r_gripper_l_finger_joint" kp="10000" name="gripper_r_gripper_l_finger_joint" forcelimited="true" forcerange="-20 20"/>
     </actuator>
 
@@ -22,32 +20,28 @@
             <body name="right_gripper" pos="0 0 0.095">
                 <inertial pos="0 0 0" mass="0.0001" diaginertia="0 0 0" />
             </body>
-            <!-- TODO: play with y-pos here (make it positive, was -0.0015) to align to center -->
-            <body name="r_gripper_l_finger" pos="0 -0.009 0.02">
+            <body name="r_gripper_l_finger" pos="0 -0.0015 0.02">
                 <inertial pos="0 0 0" quat="0 0 0 -1" mass="0.02" diaginertia="0.01 0.01 0.01" />
-                <!-- TODO: play with range -->
-                <joint name="r_gripper_l_finger_joint" pos="0 0 0" axis="0 1 0" type="slide" limited="true" range="-0.01 0.020833" damping="0.7"/>
+                <joint name="r_gripper_l_finger_joint" pos="0 0 0" axis="0 1 0" type="slide" limited="true" range="-0.008 0.020833" damping="0.7"/>
                 <geom quat="0 0 0 -1" type="mesh" contype="0" conaffinity="0" group="1" mesh="standard_narrow" />
-                <geom size="0.005 0.00675 0.0375" pos="0 0.01725 0.04" quat="0 0 0 -1" type="box" />
-                <geom size="0.005 0.025 0.0085" pos="-0.005 -0.003 0.0083" quat="0 0 0 -1" type="box" />
+                <geom size="0.005 0.00675 0.0375" pos="0 0.01725 0.04" quat="0 0 0 -1" type="box"  conaffinity="1" contype="0"/>
+                <geom size="0.005 0.025 0.0085" pos="-0.005 -0.003 0.0083" quat="0 0 0 -1" type="box"  conaffinity="1" contype="0"/>
                 <body name="r_gripper_l_finger_tip" pos="0 0.01725 0.075">
                     <inertial pos="0 0 0" quat="0 0 0 1" mass="0.01" diaginertia="0.01 0.01 0.01" />
-                    <geom quat="0 0 0 1" type="mesh" contype="0" conaffinity="0" group="1" mesh="half_round_tip" />
-                    <geom size="0.004 0.004 0.0185" pos="0 -0.0045 -0.015" quat="0 0 0 1" type="box" />
+                    <geom quat="0 0 0 1" type="mesh" contype="0" conaffinity="0" group="1" mesh="half_round_tip"/>
+                    <geom size="0.004 0.004 0.0185" pos="0 -0.0045 -0.015" quat="0 0 0 1" type="box"  conaffinity="1" contype="0"/>
                 </body>
             </body>
-            <!-- TODO: play with y-pos here (make it negative, was 0.0015) to align to center (-0.1)? -->
-            <body name="r_gripper_r_finger" pos="0 -0.006 0.02">
+            <body name="r_gripper_r_finger" pos="0 0.0015 0.02">
                 <inertial pos="0 0 0" mass="0.02" diaginertia="0.01 0.01 0.01" />
-                <!-- TODO: play with range -->
-                <joint name="r_gripper_r_finger_joint" pos="0 0 0" axis="0 1 0" type="slide" limited="true" range="-0.020833 0.01" damping="0.7"/>
+                <joint name="r_gripper_r_finger_joint" pos="0 0 0" axis="0 1 0" type="slide" limited="true" range="-0.020833 0.008" damping="0.7"/>
                 <geom type="mesh" contype="0" conaffinity="0" group="1" mesh="standard_narrow" />
-                <geom size="0.005 0.00675 0.0375" pos="0 -0.01725 0.04" type="box" />
-                <geom size="0.005 0.025 0.0085" pos="0.005 0.003 0.0083" type="box" />
+                <geom size="0.005 0.00675 0.0375" pos="0 -0.01725 0.04" type="box"  conaffinity="1" contype="0"/>
+                <geom size="0.005 0.025 0.0085" pos="0.005 0.003 0.0083" type="box"  conaffinity="1" contype="0"/>
                 <body name="r_gripper_r_finger_tip" pos="0 -0.01725 0.075">
                     <inertial pos="0 0 0" mass="0.01" diaginertia="0.01 0.01 0.01" />
-                    <geom type="mesh" contype="0" conaffinity="0" group="1" mesh="half_round_tip" />
-                    <geom size="0.004 0.004 0.0185" pos="0 0.0045 -0.015" type="box" />
+                    <geom type="mesh" contype="0" conaffinity="0" group="1" mesh="half_round_tip"/>
+                    <geom size="0.004 0.004 0.0185" pos="0 0.0045 -0.015" type="box"  conaffinity="1" contype="0"/>
                 </body>
             </body>
         </body>

MujocoManip/model/assets/object/object_cylinder.xml

@@ -0,0 +1,17 @@
+<!-- sample of a general object defined in mujoco -->
+<mujoco model="cylinder">
+    <worldbody>
+        <body name="collision">
+            <!-- Size = Radius, Half-height -->
+            <geom pos="0 0 0" size="0.05 0.08" type="cylinder" rgba="1 0 0 1"/>
+        </body>
+        <body name="visual">
+            <geom pos="0 0 0" size="0.05 0.08" type="cylinder" rgba="1 0 0 1"  conaffinity="0" contype="0" mass="0.0001" group="1"/>
+            <geom pos="0 0 0" size="0.05 0.08" type="cylinder" rgba="1 0 0 1"  conaffinity="0" contype="0" mass="0.0001" group="0"/>
+        </body>
+        <!-- These sites are for stacking objects, for aligning objects on a table for example. -->
+        <site pos="0 0 -0.08" name="bottom_site"/>
+        <site pos="0 0 0.08" name="top_site"/>
+        <site pos="0.05 0 0" name="horizontal_radius_site"/>
+    </worldbody>
+</mujoco>

MujocoManip/model/mujoco_object.py

@@ -112,6 +112,10 @@ class DefaultBallObject(MujocoXMLObject):
     def __init__(self):
         super().__init__(xml_path_completion('object/object_ball.xml'))
 
+class DefautCylinderObject(MujocoXMLObject):
+    def __init__(self):
+        super().__init__(xml_path_completion('object/object_cylinder.xml'))
+
 class MujocoGeneratedObject(MujocoObject):
     """
         Base class for all programmatically generated mujoco object
@@ -130,6 +134,7 @@ class BoxObject(MujocoGeneratedObject):
     # TODO: friction, etc
     def __init__(self, size, rgba):
         super().__init__()
+        assert(len(size) == 3)
         self.size = np.array(size)
         self.rgba = np.array(rgba)
 
@@ -167,6 +172,51 @@ def get_visual(self):
         body.append(ET.Element('geom', attrib=template_gp1))
         return body
 
+class CylinderObject(MujocoGeneratedObject):
+    """
+        An object that is a box
+    """
+    # TODO: friction, etc
+    def __init__(self, size, rgba):
+        super().__init__()
+        assert(len(size) == 2)
+        self.size = np.array(size)
+        self.rgba = np.array(rgba)
+
+    def get_bottom_offset(self):
+        return np.array([0, 0, -1 * self.size[1]])
+
+    def get_top_offset(self):
+        return np.array([0, 0, self.size[1]])
+
+    def get_horizontal_radius(self):
+        return self.size[0]
+
+    # returns a copy, Returns xml body node
+    def get_collision(self):
+        body = ET.Element('body')
+        template = self.get_collision_attrib_template()
+        template['type'] = 'cylinder'
+        template['rgba'] = array_to_string(self.rgba)
+        template['size'] = array_to_string(self.size)
+        body.append(ET.Element('geom', attrib=template))
+        return body
+
+    # returns a copy, Returns xml body node
+    def get_visual(self):
+        body = ET.Element('body')
+        template = self.get_visual_attrib_template()
+        template['type'] = 'cylinder'
+        template['rgba'] = array_to_string(self.rgba)
+        # shrink so that we don't see flickering when showing both visual and collision
+        template['size'] = array_to_string(self.size * visual_size_shrink_ratio) 
+        template['group'] = '0'
+        body.append(ET.Element('geom', attrib=template))
+        template_gp1 = copy.deepcopy(template)
+        template_gp1['group'] = '1'
+        body.append(ET.Element('geom', attrib=template_gp1))
+        return body
+
 class RandomBoxObject(BoxObject):
     """
         A random box
@@ -177,3 +227,16 @@ def __init__(self, size_max=[0.07, 0.07, 0.07], size_min=[0.03, 0.03, 0.03], see
         size = np.array([np.random.uniform(size_min[i], size_max[i]) for i in range(3)])
         rgba = np.array([np.random.uniform(0, 1) for i in range(3)] + [1])
         super().__init__(size, rgba)
+
+class RandomCylinderObject(CylinderObject):
+    """
+        A random box
+    """
+    def __init__(self, size_max=[0.07, 0.07], size_min=[0.03, 0.03], seed=None):
+        if seed is not None:
+            np.random.seed(seed)
+        size = np.array([np.random.uniform(size_min[i], size_max[i]) for i in range(2)])
+        rgba = np.array([np.random.uniform(0, 1) for i in range(3)] + [1])
+        super().__init__(size, rgba)
+
+

MujocoManip/model/robot.py

@@ -77,4 +77,6 @@ def joints(self):
 
     @property
     def rest_pos(self):
-        return [0, -1.18, 0.00, 2.18, 0.00, 0.57, 3.3161]
+        return [0, -1.18, 0.00, 2.18, 0.00, 0.57, 3.3161]
+
+