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]))
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.)
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_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.)
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 __init__(self,
players,
arena,
ball=None,
initializer=None,
observables=None,
disable_walker_contacts=False,
nconmax_per_player=200,
njmax_per_player=200,
control_timestep=0.025):
"""Construct an instance of soccer.Task.
This task implements the high-level game logic of multi-agent MuJoCo soccer.
Args:
players: a sequence of `soccer.Player` instances, representing
participants to the game from both teams.
arena: an instance of `soccer.Pitch`, implementing the physical geoms and
the sensors associated with the pitch.
ball: optional instance of `soccer.SoccerBall`, implementing the physical
geoms and sensors associated with the soccer ball. If None, defaults to
using `soccer_ball.SoccerBall()`.
initializer: optional instance of `soccer.Initializer` that initializes
the task at the start of each episode. If None, defaults to
`initializers.UniformInitializer()`.
observables: optional instance of `soccer.Observables` that adds
observables for each player. If None, defaults to
`observables.CoreObservables()`.
disable_walker_contacts: if `True`, disable physical contacts between
players.
nconmax_per_player: allocated maximum number of contacts per player. It
may be necessary to increase this value if you encounter errors due to
`mjWARN_CONTACTFULL`.
njmax_per_player: allocated maximum number of scalar constraints per
player. It may be necessary to increase this value if you encounter
errors due to `mjWARN_CNSTRFULL`.
control_timestep: control timestep of the agent.
"""
self.arena = arena
self.players = players
self._initializer = initializer or initializers.UniformInitializer()
self._observables = observables or observables_lib.CoreObservables()
if disable_walker_contacts:
_disable_geom_contacts([p.walker for p in self.players])
# Create ball and attach ball to arena.
self.ball = ball or soccer_ball.SoccerBall()
self.arena.add_free_entity(self.ball)
self.arena.register_ball(self.ball)
# Register soccer ball contact tracking for players.
for player in self.players:
player.walker.create_root_joints(self.arena.attach(player.walker))
self.ball.register_player(player)
# Add per-walkers observables.
self._observables(self, player)
self.set_timesteps(physics_timestep=0.005,
control_timestep=control_timestep)
self.root_entity.mjcf_model.size.nconmax = nconmax_per_player * len(
players)
self.root_entity.mjcf_model.size.njmax = njmax_per_player * len(
players)
