Skip to content
Permalink
Browse files

Ensure soccer initializer sets position of walker correctly regardles…

…s of walker initializer.

This replaces the assumption that the walker pose initializer sets the (x, y) position to (0., 0.).

Also added unit tests for UniformInitializer.

PiperOrigin-RevId: 250300853
  • Loading branch information...
DeepMind authored and alimuldal committed May 28, 2019
1 parent 7a36377 commit c0ff41fb5ddf12fe283311f9d544839d02a15b17
Showing with 134 additions and 9 deletions.
  1. +12 −7 dm_control/locomotion/soccer/initializers.py
  2. +122 −2 dm_control/locomotion/soccer/task_test.py
@@ -47,19 +47,24 @@ def __init__(self, spawn_ratio=_SPAWN_RATIO, init_ball_z=_INIT_BALL_Z):
def _initialize_ball(self, ball, spawn_range, physics, random_state):
x, y = random_state.uniform(-spawn_range, spawn_range)
ball.set_pose(physics, [x, y, self._init_ball_z])
ball.set_velocity(
physics,
velocity=np.asarray([0., 0., 0.]),
angular_velocity=np.zeros(3))
# Note: this method is not always called immediately after `physics.reset()`
# so we need to explicitly zero out the velocity.
ball.set_velocity(physics, velocity=0., angular_velocity=0.)

def _initialize_walker(self, walker, spawn_range, physics, random_state):
"""Uniformly initialize walker in spawn_range."""
walker.reinitialize_pose(physics, random_state)
x, y = random_state.uniform(-spawn_range, spawn_range)
(_, _, z), quat = walker.get_pose(physics)
walker.set_pose(physics, [x, y, z], quat)
rotation = random_state.uniform(-np.pi, np.pi)
quat = [np.cos(rotation / 2), 0, 0, np.sin(rotation / 2)]
walker.shift_pose(physics, [x, y, 0.], quat)
walker.set_velocity(
physics, velocity=np.zeros(3), angular_velocity=np.zeros(3))
walker.shift_pose(physics, quaternion=quat)
# TODO(b/132759890): `walker.set_velocity` has no effect for walkers without
# freejoints, such as `BoxHead`.
# Note: this method is not always called immediately after `physics.reset()`
# so we need to explicitly zero out the velocity.
walker.set_velocity(physics, velocity=0., angular_velocity=0.)

def __call__(self, task, physics, random_state):
spawn_range = np.asarray(task.arena.size) * self._spawn_ratio
@@ -19,17 +19,20 @@
from __future__ import division
from __future__ import print_function

import unittest
# Internal dependencies.

from absl.testing import absltest
from absl.testing import parameterized
from dm_control import composer
from dm_control import mjcf
from dm_control.locomotion import soccer
from dm_control.locomotion.soccer import initializers
from dm_control.mujoco.wrapper import mjbindings
import numpy as np
from six.moves import range
from six.moves import zip


RGBA_BLUE = [.1, .1, .8, 1.]
RGBA_RED = [.8, .1, .1, 1.]

@@ -58,14 +61,17 @@ def _away_team(team_size):
return _team_players(team_size, soccer.Team.AWAY, "away", RGBA_RED)


def _env(players, disable_walker_contacts=True, observables=None):
def _env(players, disable_walker_contacts=True, observables=None,
random_state=42, **task_kwargs):
return composer.Environment(
task=soccer.Task(
players=players,
arena=soccer.Pitch((20, 15)),
observables=observables,
disable_walker_contacts=disable_walker_contacts,
**task_kwargs
),
random_state=random_state,
time_limit=1)


@@ -346,5 +352,119 @@ def _initial_configuration(physics, unused_random_state):
self.assertEqual(timestep.discount, expected_terminal_discount)


class UniformInitializerTest(parameterized.TestCase):

@parameterized.parameters([0.3, 0.7])
def test_walker_position(self, spawn_ratio):
initializer = initializers.UniformInitializer(spawn_ratio=spawn_ratio)
env = _env(_home_team(2) + _away_team(2), initializer=initializer)
root_bodies = [p.walker.root_body for p in env.task.players]
xy_bounds = np.asarray(env.task.arena.size) * spawn_ratio
env.reset()
xy = env.physics.bind(root_bodies).xpos[:, :2].copy()
with self.subTest("X and Y positions within bounds"):
if np.any(abs(xy) > xy_bounds):
self.fail("Walker(s) spawned out of bounds. Expected abs(xy) "
"<= {}, got:\n{}".format(xy_bounds, xy))
env.reset()
xy2 = env.physics.bind(root_bodies).xpos[:, :2].copy()
with self.subTest("X and Y positions change after reset"):
if np.any(xy == xy2):
self.fail("Walker(s) have the same X and/or Y coordinates before and "
"after reset. Before: {}, after: {}.".format(xy, xy2))

def test_walker_rotation(self):
initializer = initializers.UniformInitializer()
env = _env(_home_team(2) + _away_team(2), initializer=initializer)

def quats_to_eulers(quats):
eulers = np.empty((len(quats), 3), dtype=np.double)
dt = 1.
for i, quat in enumerate(quats):
mjbindings.mjlib.mju_quat2Vel(eulers[i], quat, dt)
return eulers

# TODO(b/132671988): Switch to using `get_pose` to get the quaternion once
# `BoxHead.get_pose` and `BoxHead.set_pose` are
# implemented in a consistent way.
def get_quat(walker):
return env.physics.bind(walker.root_body).xquat

env.reset()
quats = [get_quat(p.walker) for p in env.task.players]
eulers = quats_to_eulers(quats)
with self.subTest("Rotation is about the Z-axis only"):
np.testing.assert_array_equal(eulers[:, :2], 0.)

env.reset()
quats2 = [get_quat(p.walker) for p in env.task.players]
eulers2 = quats_to_eulers(quats2)
with self.subTest("Rotation about Z changes after reset"):
if np.any(eulers[:, 2] == eulers2[:, 2]):
self.fail("Walker(s) have the same rotation about Z before and "
"after reset. Before: {}, after: {}."
.format(eulers[:, 2], eulers2[:, 2]))

# TODO(b/132759890): Remove `expectedFailure` decorator once `set_velocity`
# works correctly for the `BoxHead` walker.
@unittest.expectedFailure
def test_walker_velocity(self):
initializer = initializers.UniformInitializer()
env = _env(_home_team(2) + _away_team(2), initializer=initializer)
root_joints = []
non_root_joints = []
for player in env.task.players:
attachment_frame = mjcf.get_attachment_frame(player.walker.mjcf_model)
root_joints.extend(
attachment_frame.find_all("joint", immediate_children_only=True))
non_root_joints.extend(player.walker.mjcf_model.find_all("joint"))
# Assign a non-zero sentinel value to the velocities of all root and
# non-root joints.
sentinel_velocity = 3.14
env.physics.bind(root_joints + non_root_joints).qvel = sentinel_velocity
# The initializer should zero the velocities of the root joints, but not the
# non-root joints.
initializer(env.task, env.physics, env.random_state)
np.testing.assert_array_equal(env.physics.bind(non_root_joints).qvel,
sentinel_velocity)
np.testing.assert_array_equal(env.physics.bind(root_joints).qvel, 0.)

@parameterized.parameters([
dict(spawn_ratio=0.3, init_ball_z=0.4),
dict(spawn_ratio=0.5, init_ball_z=0.6),
])
def test_ball_position(self, spawn_ratio, init_ball_z):
initializer = initializers.UniformInitializer(
spawn_ratio=spawn_ratio, init_ball_z=init_ball_z)
env = _env(_home_team(2) + _away_team(2), initializer=initializer)
xy_bounds = np.asarray(env.task.arena.size) * spawn_ratio
env.reset()
position, _ = env.task.ball.get_pose(env.physics)
xyz = position.copy()
with self.subTest("X and Y positions within bounds"):
if np.any(abs(xyz[:2]) > xy_bounds):
self.fail("Ball spawned out of bounds. Expected abs(xy) "
"<= {}, got:\n{}".format(xy_bounds, xyz[:2]))
with self.subTest("Z position equal to `init_ball_z`"):
self.assertEqual(xyz[2], init_ball_z)
env.reset()
position, _ = env.task.ball.get_pose(env.physics)
xyz2 = position.copy()
with self.subTest("X and Y positions change after reset"):
if np.any(xyz[:2] == xyz2[:2]):
self.fail("Ball has the same XY position before and after reset. "
"Before: {}, after: {}.".format(xyz[:2], xyz2[:2]))

def test_ball_velocity(self):
initializer = initializers.UniformInitializer()
env = _env(_home_team(1) + _away_team(1), initializer=initializer)
ball_root_joint = mjcf.get_frame_freejoint(env.task.ball.mjcf_model)
# Set the velocities of the ball root joint to a non-zero sentinel value.
env.physics.bind(ball_root_joint).qvel = 3.14
initializer(env.task, env.physics, env.random_state)
# The initializer should set the ball velocity to zero.
ball_velocity = env.physics.bind(ball_root_joint).qvel
np.testing.assert_array_equal(ball_velocity, 0.)

if __name__ == "__main__":
absltest.main()

0 comments on commit c0ff41f

Please sign in to comment.
You can’t perform that action at this time.