Add egocentric observations of others' end_effectors.

PiperOrigin-RevId: 300648743
Change-Id: Idce50d18cb2b0bfece620ccd08c7f057cd1db335

dm_control/locomotion/soccer/observables.py

@@ -114,14 +114,24 @@ def _add_player_observables_on_other(self, player, other, prefix):
     # Origin callable in xpos, xvel for `player`.
     xpos_xyz_callable = lambda p: p.bind(player.walker.root_body).xpos
     xvel_xyz_callable = lambda p: p.bind(player.walker.root_body).cvel[3:]
-
     # Egocentric observation of other's position, orientation and
     # linear velocities.
     def _cvel_observation(physics, other=other):
       # Velocitmeter reads in local frame but we need world frame observable
       # for egocentric transformation.
       return physics.bind(other.walker.root_body).cvel[3:]
 
+    def _egocentric_end_effectors_xpos(physics, other=other):
+      origin_xpos = xpos_xyz_callable(physics)
+      egocentric_end_effectors_xpos = []
+      for end_effector_body in other.walker.end_effectors:
+        xpos = physics.bind(end_effector_body).xpos
+        delta = xpos - origin_xpos
+        ego_xpos = player.walker.transform_vec_to_egocentric_frame(
+            physics, delta)
+        egocentric_end_effectors_xpos.append(ego_xpos)
+      return np.concatenate(egocentric_end_effectors_xpos)
+
     player.walker.observables.add_egocentric_vector(
         '{}_ego_linear_velocity'.format(prefix),
         base_observable.Generic(_cvel_observation),
@@ -134,6 +144,11 @@ def _cvel_observation(physics, other=other):
         '{}_ego_orientation'.format(prefix),
         other.walker.observables.orientation)
 
+    # Adds end effectors of the other agents in the player's egocentric frame.
+    player.walker.observables.add_observable(
+        '{}_ego_end_effectors_pos'.format(prefix),
+        base_observable.Generic(_egocentric_end_effectors_xpos))
+
     # Adds end effectors of the other agents in the other's egocentric frame.
     # A is seeing B's hand extended to B's right.
     player.walker.observables.add_observable(

dm_control/locomotion/soccer/task_test.py

@@ -97,14 +97,14 @@ def _assert_all_count_equal(self, list_of_lists):
       self.assertCountEqual(first, other)
 
   @parameterized.named_parameters(
-      ("1vs1_core", 1, "core", 32, True),
-      ("2vs2_core", 2, "core", 40, True),
-      ("1vs1_interception", 1, "core_interception", 40, True),
-      ("2vs2_interception", 2, "core_interception", 48, True),
-      ("1vs1_core_contact", 1, "core", 32, False),
-      ("2vs2_core_contact", 2, "core", 40, False),
-      ("1vs1_interception_contact", 1, "core_interception", 40, False),
-      ("2vs2_interception_contact", 2, "core_interception", 48, False),
+      ("1vs1_core", 1, "core", 33, True),
+      ("2vs2_core", 2, "core", 43, True),
+      ("1vs1_interception", 1, "core_interception", 41, True),
+      ("2vs2_interception", 2, "core_interception", 51, True),
+      ("1vs1_core_contact", 1, "core", 33, False),
+      ("2vs2_core_contact", 2, "core", 43, False),
+      ("1vs1_interception_contact", 1, "core_interception", 41, False),
+      ("2vs2_interception_contact", 2, "core_interception", 51, False),
   )
   def test_step_environment(self, team_size, observables_adder, num_obs,
                             disable_walker_contacts):
@@ -130,11 +130,11 @@ def test_step_environment(self, team_size, observables_adder, num_obs,
 
   # TODO(b/124848293): consolidate environment stepping loop for task tests.
   @parameterized.named_parameters(
-      ("1vs2", 1, 2, 36),
-      ("2vs1", 2, 1, 36),
-      ("3vs0", 3, 0, 36),
-      ("0vs2", 0, 2, 32),
-      ("2vs2", 2, 2, 40),
+      ("1vs2", 1, 2, 38),
+      ("2vs1", 2, 1, 38),
+      ("3vs0", 3, 0, 38),
+      ("0vs2", 0, 2, 33),
+      ("2vs2", 2, 2, 43),
       ("0vs0", 0, 0, None),
   )
   def test_num_players(self, home_size, away_size, num_observations):

setup.py

@@ -166,7 +166,7 @@ def find_data_files(package_dir, patterns):
 
 setup(
     name='dm_control',
-    version='0.0.300565743',
+    version='0.0.300648743',
     description='Continuous control environments and MuJoCo Python bindings.',
     author='DeepMind',
     license='Apache License, Version 2.0',