Adding a movable head to the BoxHead walker with an attached camera to allow for egocentric vision.

DeepMind · copybara-github · commit 308b73c40994 · 2020-03-23T03:31:58.000-07:00
PiperOrigin-RevId: 302398491
Change-Id: I0da293b76507a1841d7943327933f14e18499821
diff --git a/dm_control/locomotion/soccer/assets/boxhead/boxhead.xml b/dm_control/locomotion/soccer/assets/boxhead/boxhead.xml
@@ -6,46 +6,39 @@
       <joint pos='0 0 0' limited='false' armature='0' damping='0' stiffness='0'/>
     </default>
   </default>
-
   <worldbody>
     <body pos='0 0 1.0' quat='1 0 0 -1' name='head_body'>
       <camera name="float_far"  pos="-4 0 2" xyaxes="0 -1 0 .5 0 1" fovy="90"/>
-      <body name='egocentric_camera' pos='0.1 0 0.2' xyaxes='0 -1 0 0.09950371902 0 0.99503719021'>
+      <body name='egocentric_camera' pos='0 0 0.4' xyaxes='0 -1 0 0 0 1'>
         <inertial pos='0 0 0' mass='1' diaginertia='1 1 1'/>
-        <joint name='camera_yaw' type='hinge' axis='0 1 0' limited='true' range='-0.5 0.5'/>
-        <joint name='camera_pitch' type='hinge' axis='1 0 0' limited='true' range='-0.35 0.35'/>
-        <geom type='cylinder' name='eye_l' pos='-.14 -0.2 -0.171' size="0.06 0.05" rgba='1 1 1 1' mass='0' contype="0" conaffinity="0"/>
-        <geom type='cylinder' name='eye_r' pos='.14 -0.2 -0.171' size="0.06 0.05" rgba='1 1 1 1' mass='0' contype="0" conaffinity="0"/>
+        <joint name='camera_yaw' type='hinge' axis='0 1 0' limited='false'/>
+        <joint name='camera_pitch' type='hinge' axis='1 0 0' limited='true' range='-0.35 0.00'/>
+        <geom type='cylinder' name='eye_l' pos='-.14 -0.4 -0.171' size="0.06 0.2" rgba='1 1 1 1' mass='0' contype="0" conaffinity="0"/>
+        <geom type='cylinder' name='eye_r' pos='.14 -0.4 -0.171' size="0.06 0.2" rgba='1 1 1 1' mass='0' contype="0" conaffinity="0"/>
         <camera name='egocentric' fovy='80'/>
       </body>
       <camera name='tracking' pos='1.5 1.5 .7' xyaxes='-1 1 0 -1 -1 4' mode='trackcom'/>
       <joint name='steer' type='hinge' axis='0 0 -1' class="root" damping="4"/>
-
       <geom type='box' group='5' name='top_down_cam_box' pos='0 0 2' size="1. 1. 0.1" rgba='.1 .1 .1 1' mass='0' contype="0" conaffinity="0"/>
       <geom type='box' name='head' pos='0 0 0' size="0.2 0.2 0.1" rgba='.1 .1 .1 1' mass='20' contype="1" conaffinity="1"/>
       <site name='torso_site' pos='0 0 0' size='0.05' rgba='1 0 0 1' group='4'/>
-
       <body pos='0 0 0.05' name='torso'>
         <!-- wheel -->
         <joint name='kick' type='slide' damping='100' limited='true' range='0 .4' axis="0 0 -1" stiffness="0" solimplimit=".95 .99 .001" springref="-.4"/>
         <geom name='arm_l' type='capsule' fromto='0  .35 -.7 .03  .8 -.7' size='.1' rgba='.76 .76 .76 1' density='10'/>
         <geom name='arm_r' type='capsule' fromto='0 -.35 -.7 .03 -.8 -.7' size='.1' rgba='.76 .76 .76 1' density='10'/>
-
         <body name='ball' pos='0 0 -0.7'>
           <joint name='roll' type='hinge' axis='0 1 0' damping='2'/>
           <geom name='shell' type='sphere' size='.35' density='100' condim='4' friction='.5 .02 .02'/>
         </body>
-
       </body>
     </body>
   </worldbody>
-
   <sensor>
     <velocimeter site="torso_site" name="sensor_torso_vel"/>
     <gyro site="torso_site" name="sensor_torso_gyro"/>
     <accelerometer site="torso_site" name="sensor_torso_accel"/>
   </sensor>
-
   <actuator>
     <general name='roll' joint='roll' ctrlrange='-1 1' ctrllimited='true' gear='-30' biasprm="0" biastype="affine"/>
     <motor name='steer' joint='steer' ctrlrange='-1 1' ctrllimited='true' gear='55'/>
diff --git a/dm_control/locomotion/soccer/boxhead.py b/dm_control/locomotion/soccer/boxhead.py
@@ -24,6 +24,7 @@
 from dm_control import composer
 from dm_control import mjcf
 from dm_control.composer.observation import observable
+from dm_control.locomotion.walkers import initializers
 from dm_control.locomotion.walkers import legacy_base
 import numpy as np
 from PIL import Image
@@ -104,6 +105,7 @@ def _asset_png_with_background_rgba_bytes(asset_fname, background_rgba):
 
 
 class BoxHeadObservables(legacy_base.WalkerObservables):
+  """BoxHead observables with low-res camera and modulo'd rotational joints."""
 
   def __init__(self, entity, camera_resolution):
     self._camera_resolution = camera_resolution
@@ -115,6 +117,43 @@ def egocentric_camera(self):
     return observable.MJCFCamera(self._entity.egocentric_camera,
                                  width=width, height=height)
 
+  @property
+  def proprioception(self):
+    proprioception = super(BoxHeadObservables, self).proprioception
+    if self._entity.observable_camera_joints:
+      return proprioception + [self.camera_joints_pos, self.camera_joints_vel]
+    return proprioception
+
+  @composer.observable
+  def camera_joints_pos(self):
+
+    def _sin(value, random_state):
+      del random_state
+      return np.sin(value)
+
+    def _cos(value, random_state):
+      del random_state
+      return np.cos(value)
+
+    sin_rotation_joints = observable.MJCFFeature(
+        'qpos', self._entity.observable_camera_joints, corruptor=_sin)
+
+    cos_rotation_joints = observable.MJCFFeature(
+        'qpos', self._entity.observable_camera_joints, corruptor=_cos)
+
+    def _camera_joints(physics):
+      return np.concatenate([
+          sin_rotation_joints(physics),
+          cos_rotation_joints(physics)
+      ], -1)
+
+    return observable.Generic(_camera_joints)
+
+  @composer.observable
+  def camera_joints_vel(self):
+    return observable.MJCFFeature(
+        'qvel', self._entity.observable_camera_joints)
+
 
 class BoxHead(legacy_base.Walker):
   """A rollable and jumpable ball with a head."""
@@ -146,7 +185,8 @@ def _build(self,
     Raises:
       ValueError: if received invalid walker_id.
     """
-    super(BoxHead, self)._build(initializer=initializer)
+    super(BoxHead, self)._build(
+        initializer=initializer or initializers.NoOpInitializer())
     xml_path = os.path.join(_ASSETS_PATH, 'boxhead.xml')
     self._mjcf_root = mjcf.from_xml_string(resources.GetResource(xml_path, 'r'))
     if name:
@@ -246,7 +286,9 @@ def set_pose(self, physics, position=None, quaternion=None):
           1 - 2 * (quaternion[2] ** 2 + quaternion[3] ** 2))
       physics.bind(self._mjcf_root.find('joint', 'steer')).qpos = z_angle
 
-  def initialize_episode(self, physics, unused_random_state):
+  def initialize_episode(self, physics, random_state):
+    self.reinitialize_pose(physics, random_state)
+
     if self._camera_control:
       _compensate_gravity(physics,
                           self._mjcf_root.find('body', 'egocentric_camera'))
@@ -279,6 +321,15 @@ def end_effectors(self):
   def observable_joints(self):
     return (self._mjcf_root.find('joint', 'kick'),)
 
+  @composer.cached_property
+  def observable_camera_joints(self):
+    if self._camera_control:
+      return (
+          self._mjcf_root.find('joint', 'camera_yaw'),
+          self._mjcf_root.find('joint', 'camera_pitch'),
+      )
+    return ()
+
   @composer.cached_property
   def egocentric_camera(self):
     return self._mjcf_root.find('camera', 'egocentric')
diff --git a/dm_control/locomotion/walkers/initializers/__init__.py b/dm_control/locomotion/walkers/initializers/__init__.py
@@ -49,7 +49,7 @@ def initialize_pose(self, physics, walker, random_state):
 
 
 class RandomlySampledInitializer(WalkerInitializer):
-  """Initializer that random selects between many initializers."""
+  """An initializer that random selects between many initializers."""
 
   def __init__(self, initializers):
     self._initializers = initializers
@@ -59,3 +59,13 @@ def initialize_pose(self, physics, walker, random_state):
     random_initalizer_idx = np.random.randint(0, self.num_initializers)
     self._initializers[random_initalizer_idx].initialize_pose(
         physics, walker, random_state)
+
+
+class NoOpInitializer(WalkerInitializer):
+  """An initializer that does nothing."""
+
+  def initialize_pose(self, physics, walker, random_state):
+    pass
+
+
+
diff --git a/setup.py b/setup.py
@@ -166,7 +166,7 @@ def find_data_files(package_dir, patterns):
 
 setup(
     name='dm_control',
-    version='0.0.301472480',
+    version='0.0.302398491',
     description='Continuous control environments and MuJoCo Python bindings.',
     author='DeepMind',
     license='Apache License, Version 2.0',
