# Copyright 2020 The dm_control Authors.

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

# ============================================================================

"""Dog Domain."""

import collections

import os

from dm_control import mujoco

from dm_control.rl import control

from dm_control.suite import base

from dm_control.suite import common

from dm_control.utils import containers

from dm_control.utils import rewards

from dm_control.utils import xml_tools

from lxml import etree

import numpy as np

from dm_control.utils import io as resources

_DEFAULT_TIME_LIMIT = 15

_CONTROL_TIMESTEP = .015

# Angle (in degrees) of local z from global z below which upright reward is 1.

_MAX_UPRIGHT_ANGLE = 30

_MIN_UPRIGHT_COSINE = np.cos(np.deg2rad(_MAX_UPRIGHT_ANGLE))

# Standing reward is 1 for body-over-foot height that is at least this fraction

# of the height at the default pose.

_STAND_HEIGHT_FRACTION = 0.9

# Torques which enforce joint range limits should stay below this value.

_EXCESSIVE_LIMIT_TORQUES = 150

# Horizontal speed above which Move reward is 1.

_WALK_SPEED = 1

_TROT_SPEED = 3

_RUN_SPEED = 9

_HINGE_TYPE = mujoco.wrapper.mjbindings.enums.mjtJoint.mjJNT_HINGE

_LIMIT_TYPE = mujoco.wrapper.mjbindings.enums.mjtConstraint.mjCNSTR_LIMIT_JOINT

SUITE = containers.TaggedTasks()

_ASSET_DIR = os.path.join(os.path.dirname(__file__), 'dog_assets')

def make_model(floor_size, remove_ball):

"""Sets floor size, removes ball and walls (Stand and Move tasks)."""

xml_string = common.read_model('dog.xml')

parser = etree.XMLParser(remove_blank_text=True)

mjcf = etree.XML(xml_string, parser)

# set floor size.

floor = xml_tools.find_element(mjcf, 'geom', 'floor')

floor.attrib['size'] = str(floor_size) + ' ' + str(floor_size) + ' .1'

if remove_ball:

# Remove ball, target and walls.

ball = xml_tools.find_element(mjcf, 'body', 'ball')

ball.getparent().remove(ball)

target = xml_tools.find_element(mjcf, 'geom', 'target')

target.getparent().remove(target)

ball_cam = xml_tools.find_element(mjcf, 'camera', 'ball')

ball_cam.getparent().remove(ball_cam)

head_cam = xml_tools.find_element(mjcf, 'camera', 'head')

head_cam.getparent().remove(head_cam)

for wall_name in ['px', 'nx', 'py', 'ny']:

wall = xml_tools.find_element(mjcf, 'geom', 'wall_' + wall_name)

wall.getparent().remove(wall)

return etree.tostring(mjcf, pretty_print=True)

def get_model_and_assets(floor_size=10, remove_ball=True):

"""Returns a tuple containing the model XML string and a dict of assets."""

assets = common.ASSETS.copy()

_, _, filenames = next(resources.WalkResources(_ASSET_DIR))

for filename in filenames:

assets[filename] = resources.GetResource(os.path.join(_ASSET_DIR, filename))

return make_model(floor_size, remove_ball), assets

@SUITE.add('no_reward_visualization')

def stand(time_limit=_DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):

"""Returns the Stand task."""

floor_size = _WALK_SPEED * _DEFAULT_TIME_LIMIT

physics = Physics.from_xml_string(*get_model_and_assets(floor_size))

task = Stand(random=random)

environment_kwargs = environment_kwargs or {}

return control.Environment(physics, task, time_limit=time_limit,

control_timestep=_CONTROL_TIMESTEP,

**environment_kwargs)

@SUITE.add('no_reward_visualization')

def walk(time_limit=_DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):

"""Returns the Walk task."""

move_speed = _WALK_SPEED

floor_size = move_speed * _DEFAULT_TIME_LIMIT

physics = Physics.from_xml_string(*get_model_and_assets(floor_size))

task = Move(move_speed=move_speed, random=random)

environment_kwargs = environment_kwargs or {}

return control.Environment(physics, task, time_limit=time_limit,

control_timestep=_CONTROL_TIMESTEP,

**environment_kwargs)

@SUITE.add('no_reward_visualization')

def trot(time_limit=_DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):

"""Returns the Trot task."""

move_speed = _TROT_SPEED

floor_size = move_speed * _DEFAULT_TIME_LIMIT

physics = Physics.from_xml_string(*get_model_and_assets(floor_size))

task = Move(move_speed=move_speed, random=random)

environment_kwargs = environment_kwargs or {}

return control.Environment(physics, task, time_limit=time_limit,

control_timestep=_CONTROL_TIMESTEP,

**environment_kwargs)

@SUITE.add('no_reward_visualization')

def run(time_limit=_DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):

"""Returns the Run task."""

move_speed = _RUN_SPEED

floor_size = move_speed * _DEFAULT_TIME_LIMIT

physics = Physics.from_xml_string(*get_model_and_assets(floor_size))

task = Move(move_speed=move_speed, random=random)

environment_kwargs = environment_kwargs or {}

return control.Environment(physics, task, time_limit=time_limit,

control_timestep=_CONTROL_TIMESTEP,

**environment_kwargs)

@SUITE.add('no_reward_visualization', 'hard')

def fetch(time_limit=_DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):

"""Returns the Fetch task."""

physics = Physics.from_xml_string(*get_model_and_assets(remove_ball=False))

task = Fetch(random=random)

environment_kwargs = environment_kwargs or {}

return control.Environment(physics, task, time_limit=time_limit,

control_timestep=_CONTROL_TIMESTEP,

**environment_kwargs)

class Physics(mujoco.Physics):

"""Physics simulation with additional features for the Dog domain."""

def torso_pelvis_height(self):

"""Returns the height of the torso."""

return self.named.data.xpos[['torso', 'pelvis'], 'z']

def z_projection(self):

"""Returns rotation-invariant projection of local frames to the world z."""

return np.vstack((self.named.data.xmat['skull', ['zx', 'zy', 'zz']],

self.named.data.xmat['torso', ['zx', 'zy', 'zz']],

self.named.data.xmat['pelvis', ['zx', 'zy', 'zz']]))

def upright(self):

"""Returns projection from local z-axes to the z-axis of world."""

return self.z_projection()[:, 2]

def center_of_mass_velocity(self):

"""Returns the velocity of the center-of-mass."""

return self.named.data.sensordata['torso_linvel']

def torso_com_velocity(self):

"""Returns the velocity of the center-of-mass in the torso frame."""

torso_frame = self.named.data.xmat['torso'].reshape(3, 3).copy()

return self.center_of_mass_velocity().dot(torso_frame)

def com_forward_velocity(self):

"""Returns the com velocity in the torso's forward direction."""

return self.torso_com_velocity()[0]

def joint_angles(self):

"""Returns the configuration of all hinge joints (skipping free joints)."""

hinge_joints = self.model.jnt_type == _HINGE_TYPE

qpos_index = self.model.jnt_qposadr[hinge_joints]

return self.data.qpos[qpos_index].copy()

def joint_velocities(self):

"""Returns the velocity of all hinge joints (skipping free joints)."""

hinge_joints = self.model.jnt_type == _HINGE_TYPE

qvel_index = self.model.jnt_dofadr[hinge_joints]

return self.data.qvel[qvel_index].copy()

def inertial_sensors(self):

"""Returns inertial sensor readings."""

return self.named.data.sensordata[['accelerometer', 'velocimeter', 'gyro']]

def touch_sensors(self):

"""Returns touch readings."""

return self.named.data.sensordata[['palm_L', 'palm_R', 'sole_L', 'sole_R']]

def foot_forces(self):

"""Returns touch readings."""

return self.named.data.sensordata[['foot_L', 'foot_R', 'hand_L', 'hand_R']]

def ball_in_head_frame(self):

"""Returns the ball position and velocity in the frame of the head."""

head_frame = self.named.data.site_xmat['head'].reshape(3, 3)

head_pos = self.named.data.site_xpos['head']

ball_pos = self.named.data.geom_xpos['ball']

head_to_ball = ball_pos - head_pos

head_vel, _ = self.data.object_velocity('head', 'site')

ball_vel, _ = self.data.object_velocity('ball', 'geom')

head_to_ball_vel = ball_vel - head_vel

return np.hstack((head_to_ball.dot(head_frame),

head_to_ball_vel.dot(head_frame)))

def target_in_head_frame(self):

"""Returns the target position in the frame of the head."""

head_frame = self.named.data.site_xmat['head'].reshape(3, 3)

head_pos = self.named.data.site_xpos['head']

target_pos = self.named.data.geom_xpos['target']

head_to_target = target_pos - head_pos

return head_to_target.dot(head_frame)

def ball_to_mouth_distance(self):

"""Returns the distance from the ball to the mouth."""

ball_pos = self.named.data.geom_xpos['ball']

upper_bite_pos = self.named.data.site_xpos['upper_bite']

lower_bite_pos = self.named.data.site_xpos['lower_bite']

upper_dist = np.linalg.norm(ball_pos - upper_bite_pos)

lower_dist = np.linalg.norm(ball_pos - lower_bite_pos)

return 0.5*(upper_dist + lower_dist)

def ball_to_target_distance(self):

"""Returns the distance from the ball to the target."""

ball_pos, target_pos = self.named.data.geom_xpos[['ball', 'target']]

return np.linalg.norm(ball_pos - target_pos)

class Stand(base.Task):

"""A dog stand task generating upright posture."""

def __init__(self, random=None, observe_reward_factors=False):

"""Initializes an instance of `Stand`.

Args:

random: Optional, either a `numpy.random.RandomState` instance, an

integer seed for creating a new `RandomState`, or None to select a seed

automatically (default).

observe_reward_factors: Boolean, whether the factorised reward is a

key in the observation dict returned to the agent.

"""

self._observe_reward_factors = observe_reward_factors

super().__init__(random=random)

def initialize_episode(self, physics):

"""Randomizes initial root velocities and actuator states.

Args:

physics: An instance of `Physics`.

"""

physics.reset()

# Measure stand heights from default pose, above which stand reward is 1.

self._stand_height = physics.torso_pelvis_height() * _STAND_HEIGHT_FRACTION

# Measure body weight.

body_mass = physics.named.model.body_subtreemass['torso']

self._body_weight = -physics.model.opt.gravity[2] * body_mass

# Randomize horizontal orientation.

azimuth = self.random.uniform(0, 2*np.pi)

orientation = np.array((np.cos(azimuth/2), 0, 0, np.sin(azimuth/2)))

physics.named.data.qpos['root'][3:] = orientation

# Randomize root velocities in horizontal plane.

physics.data.qvel[0] = 2 * self.random.randn()

physics.data.qvel[1] = 2 * self.random.randn()

physics.data.qvel[5] = 2 * self.random.randn()

# Randomize actuator states.

assert physics.model.nu == physics.model.na

for actuator_id in range(physics.model.nu):

ctrlrange = physics.model.actuator_ctrlrange[actuator_id]

physics.data.act[actuator_id] = self.random.uniform(*ctrlrange)

def get_observation_components(self, physics):

"""Returns the observations for the Stand task."""

obs = collections.OrderedDict()

obs['joint_angles'] = physics.joint_angles()

obs['joint_velocites'] = physics.joint_velocities()

obs['torso_pelvis_height'] = physics.torso_pelvis_height()

obs['z_projection'] = physics.z_projection().flatten()

obs['torso_com_velocity'] = physics.torso_com_velocity()

obs['inertial_sensors'] = physics.inertial_sensors()

obs['foot_forces'] = physics.foot_forces()

obs['touch_sensors'] = physics.touch_sensors()

obs['actuator_state'] = physics.data.act.copy()

return obs

def get_observation(self, physics):

"""Returns the observation, possibly adding reward factors."""

obs = self.get_observation_components(physics)

if self._observe_reward_factors:

obs['reward_factors'] = self.get_reward_factors(physics)

return obs

def get_reward_factors(self, physics):

"""Returns the factorized reward."""

# Keep the torso at standing height.

torso = rewards.tolerance(physics.torso_pelvis_height()[0],

bounds=(self._stand_height[0], float('inf')),

margin=self._stand_height[0])

# Keep the pelvis at standing height.

pelvis = rewards.tolerance(physics.torso_pelvis_height()[1],

bounds=(self._stand_height[1], float('inf')),

margin=self._stand_height[1])

# Keep head, torso and pelvis upright.

upright = rewards.tolerance(physics.upright(),

bounds=(_MIN_UPRIGHT_COSINE, float('inf')),

sigmoid='linear',

margin=_MIN_UPRIGHT_COSINE+1,

value_at_margin=0)

# Reward for foot touch forces up to bodyweight.

touch = rewards.tolerance(physics.touch_sensors().sum(),

bounds=(self._body_weight, float('inf')),

margin=self._body_weight,

sigmoid='linear',

value_at_margin=0.9)

return np.hstack((torso, pelvis, upright, touch))

def get_reward(self, physics):

"""Returns the reward, product of reward factors."""

return np.prod(self.get_reward_factors(physics))

class Move(Stand):

"""A dog move task for generating locomotion."""

def __init__(self, move_speed, random, observe_reward_factors=False):

"""Initializes an instance of `Move`.

Args:

move_speed: A float. Specifies a target horizontal velocity.

random: Optional, either a `numpy.random.RandomState` instance, an

integer seed for creating a new `RandomState`, or None to select a seed

automatically (default).

observe_reward_factors: Boolean, whether the factorised reward is a

component of the observation dict.

"""

self._move_speed = move_speed

super().__init__(random, observe_reward_factors)

def get_reward_factors(self, physics):

"""Returns the factorized reward."""

standing = super().get_reward_factors(physics)

speed_margin = max(1.0, self._move_speed)

forward = rewards.tolerance(physics.com_forward_velocity(),

bounds=(self._move_speed, 2*self._move_speed),

margin=speed_margin,

value_at_margin=0,

sigmoid='linear')

forward = (4*forward + 1) / 5

return np.hstack((standing, forward))

class Fetch(Stand):

"""A dog fetch task to fetch a thrown ball."""

def __init__(self, random, observe_reward_factors=False):

"""Initializes an instance of `Move`.

Args:

random: Optional, either a `numpy.random.RandomState` instance, an

integer seed for creating a new `RandomState`, or None to select a seed

automatically (default).

observe_reward_factors: Boolean, whether the factorised reward is a

component of the observation dict.

"""

super().__init__(random, observe_reward_factors)

def initialize_episode(self, physics):

super().initialize_episode(physics)

# Set initial ball state: flying towards the center at an upward angle.

radius = 0.75 * physics.named.model.geom_size['floor', 0]

azimuth = self.random.uniform(0, 2*np.pi)

position = (radius*np.sin(azimuth), radius*np.cos(azimuth), 0.05)

physics.named.data.qpos['ball_root'][:3] = position

vertical_height = self.random.uniform(0, 3)

# Equating kinetic and potential energy: mv^2/2 = m*g*h -> v = sqrt(2gh)

gravity = -physics.model.opt.gravity[2]

vertical_velocity = np.sqrt(2 * gravity * vertical_height)

horizontal_speed = self.random.uniform(0, 5)

# Pointing towards the center, with some noise.

direction = np.array((-np.sin(azimuth) + 0.05*self.random.randn(),

-np.cos(azimuth) + 0.05*self.random.randn()))

horizontal_velocity = horizontal_speed * direction

velocity = np.hstack((horizontal_velocity, vertical_velocity))

physics.named.data.qvel['ball_root'][:3] = velocity

def get_observation_components(self, physics):

"""Returns the common observations for the Stand task."""

obs = super().get_observation_components(physics)

obs['ball_state'] = physics.ball_in_head_frame()

obs['target_position'] = physics.target_in_head_frame()

return obs

def get_reward_factors(self, physics):

"""Returns a reward to the agent."""

standing = super().get_reward_factors(physics)

# Reward for bringing mouth close to ball.

bite_radius = physics.named.model.site_size['upper_bite', 0]

bite_margin = 2

reach_ball = rewards.tolerance(physics.ball_to_mouth_distance(),

bounds=(0, bite_radius),

sigmoid='reciprocal',

margin=bite_margin)

reach_ball = (6*reach_ball + 1) / 7

# Reward for bringing the ball close to the target.

target_radius = physics.named.model.geom_size['target', 0]

bring_margin = physics.named.model.geom_size['floor', 0]

ball_near_target = rewards.tolerance(

physics.ball_to_target_distance(),

bounds=(0, target_radius),

sigmoid='reciprocal',

margin=bring_margin)

fetch_ball = (ball_near_target + 1) / 2

# Let go of the ball if it's been fetched.

if physics.ball_to_target_distance() < 2*target_radius:

reach_ball = 1

return np.hstack((standing, reach_ball, fetch_ball))