envs/custom_dmc_tasks/stickman.py

# Copyright 2017 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.
# ============================================================================
"""Stickman Domain."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import os
import numpy as np
import types

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.suite.utils import randomizers
from dm_control.utils import containers
from dm_control.utils import rewards
from dm_control.utils import io as resources
from dm_control import suite

class StickmanYogaPoses:
    lie_back   = [ -1.2 ,  0. ,  -1.57,  0, 0. , 0.0, 0, -0.,  0.0]
    lie_front  = [-1.2, -0, 1.57, 0, 0, 0, 0, 0., 0.]
    legs_up    = [ -1.24 ,  0. ,  -1.57,  1.57, 0. , 0.0,  1.57, -0.,  0.0]

    kneel      = [ -0.5 ,  0. ,  0,  0, -1.57, -0.8,  1.57, -1.57,  0.0]
    side_angle = [ -0.3 ,  0. ,  0.9,  0, 0, -0.7,  1.87, -1.07,  0.0]
    stand_up   = [-0.15, 0., 0.34, 0.74, -1.34, -0., 1.1, -0.66, -0.1]

    lean_back  = [-0.27, 0., -0.45, 0.22, -1.5, 0.86, 0.6, -0.8, -0.4]
    boat       = [ -1.04 ,  0. ,  -0.8,  1.6, 0. , 0.0, 1.6, -0.,  0.0]
    bridge     = [-1.1, 0., -2.2, -0.3, -1.5, 0., -0.3, -0.8, -0.4]

    head_stand = [-1, 0., -3, 0.6, -1, -0.3, 0.9, -0.5, 0.3]
    one_feet   = [-0.2, 0., 0, 0.7, -1.34, 0.5, 1.5, -0.6, 0.1]
    arabesque  = [-0.34, 0., 1.57, 1.57, 0, 0., 0, -0., 0.]

    # new
    high_kick = [-0.165, 3.3  , 5.55 , 1.35 ,-0, +0.5 , -0.7, 0. , 0.2,]
    splits    = [-0.7, 0., 0.5, -0.7, -1. , 0, 1.75, 0., -0.45 ]
    sit_knees = [-0.6, -0.2,  0.2, 0.95, -2.5, 0 , 0.95, -2.5, 0 ]


_DEFAULT_TIME_LIMIT = 25
_CONTROL_TIMESTEP = .025

# Minimal height of torso over foot above which stand reward is 1.
_STAND_HEIGHT = 1.15

# Horizontal speeds (meters/second) above which move reward is 1.
_WALK_SPEED = 1
_RUN_SPEED = 8

# Copied from walker:
_YOGA_HANDS_UP_HEIGHT = 1.75
_YOGA_STAND_HEIGHT = 1.0 # lower than stan height = 1.2
_YOGA_LIE_DOWN_HEIGHT = 0.1
_YOGA_LEGS_UP_HEIGHT = 1.1

_YOGA_FEET_UP_HEIGHT = 0.5
_YOGA_FEET_UP_LIE_DOWN_HEIGHT = 0.35

_YOGA_KNEE_HEIGHT = 0.25
_YOGA_KNEESTAND_HEIGHT = 0.75

_YOGA_SITTING_HEIGHT = 0.55
_YOGA_SITTING_LEGS_HEIGHT = 0.15

# speed from: https://github.com/rll-research/url_benchmark/blob/710c3eb/custom_dmc_tasks/py
_SPIN_SPEED = 5.0
#
_PUNCH_SPEED = 5.0
_PUNCH_DIST = 0.29


SUITE = containers.TaggedTasks()

def make(task,
         task_kwargs=None,
         environment_kwargs=None,
         visualize_reward=False):
    task_kwargs = task_kwargs or {}
    if environment_kwargs is not None:
        task_kwargs = task_kwargs.copy()
        task_kwargs['environment_kwargs'] = environment_kwargs
    env = SUITE[task](**task_kwargs)
    env.task.visualize_reward = visualize_reward
    return env

def get_model_and_assets():
    """Returns a tuple containing the model XML string and a dict of assets."""
    root_dir = os.path.dirname(os.path.dirname(__file__))
    xml = resources.GetResource(os.path.join(root_dir, 'custom_dmc_tasks', 'stickman.xml'))
    return xml, common.ASSETS

@SUITE.add('custom')
def hands_up(time_limit=_DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
  """Returns the hands_up task."""
  physics = Physics.from_xml_string(*get_model_and_assets())
  task = Stickman(goal='hands_up', random=random)
  environment_kwargs = environment_kwargs or {}
  return control.Environment(
      physics, task, time_limit=time_limit, control_timestep=_CONTROL_TIMESTEP,
      **environment_kwargs)


@SUITE.add('custom')
def boxing(time_limit=_DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
  """Returns the boxing task."""
  physics = Physics.from_xml_string(*get_model_and_assets())
  task = Stickman(goal='boxing', random=random)
  environment_kwargs = environment_kwargs or {}
  return control.Environment(
      physics, task, time_limit=time_limit, control_timestep=_CONTROL_TIMESTEP,
      **environment_kwargs)

@SUITE.add('custom')
def arabesque(time_limit=_DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
  """Returns the Arabesque task."""
  physics = Physics.from_xml_string(*get_model_and_assets())
  task = Stickman(goal='arabesque', random=random)
  environment_kwargs = environment_kwargs or {}
  return control.Environment(
      physics, task, time_limit=time_limit, control_timestep=_CONTROL_TIMESTEP,
      **environment_kwargs)


@SUITE.add('custom')
def lying_down(time_limit=_DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
  """Returns the Lie Down task."""
  physics = Physics.from_xml_string(*get_model_and_assets())
  task = Stickman(goal='lying_down', random=random)
  environment_kwargs = environment_kwargs or {}
  return control.Environment(
      physics, task, time_limit=time_limit, control_timestep=_CONTROL_TIMESTEP,
      **environment_kwargs)


@SUITE.add('custom')
def legs_up(time_limit=_DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
  """Returns the Legs Up task."""
  physics = Physics.from_xml_string(*get_model_and_assets())
  task = Stickman(goal='legs_up', random=random)
  environment_kwargs = environment_kwargs or {}
  return control.Environment(
      physics, task, time_limit=time_limit, control_timestep=_CONTROL_TIMESTEP,
      **environment_kwargs)

@SUITE.add('custom')
def high_kick(time_limit=_DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
  """Returns the High Kick task."""
  physics = Physics.from_xml_string(*get_model_and_assets())
  task = Stickman(goal='high_kick', random=random)
  environment_kwargs = environment_kwargs or {}
  return control.Environment(
      physics, task, time_limit=time_limit, control_timestep=_CONTROL_TIMESTEP,
      **environment_kwargs)

@SUITE.add('custom')
def one_foot(time_limit=_DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
  """Returns the High Kick task."""
  physics = Physics.from_xml_string(*get_model_and_assets())
  task = Stickman(goal='one_foot', random=random)
  environment_kwargs = environment_kwargs or {}
  return control.Environment(
      physics, task, time_limit=time_limit, control_timestep=_CONTROL_TIMESTEP,
      **environment_kwargs)

@SUITE.add('custom')
def lunge_pose(time_limit=_DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
  """Returns the High Kick task."""
  physics = Physics.from_xml_string(*get_model_and_assets())
  task = Stickman(goal='lunge_pose', random=random)
  environment_kwargs = environment_kwargs or {}
  return control.Environment(
      physics, task, time_limit=time_limit, control_timestep=_CONTROL_TIMESTEP,
      **environment_kwargs)

@SUITE.add('custom')
def sit_knees(time_limit=_DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
  """Returns the High Kick task."""
  physics = Physics.from_xml_string(*get_model_and_assets())
  task = Stickman(goal='sit_knees', random=random)
  environment_kwargs = environment_kwargs or {}
  return control.Environment(
      physics, task, time_limit=time_limit, control_timestep=_CONTROL_TIMESTEP,
      **environment_kwargs)

@SUITE.add('custom')
def headstand(time_limit=_DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
  """Returns the Headstand task."""
  physics = Physics.from_xml_string(*get_model_and_assets())
  task = Stickman(goal='flip', move_speed=0, random=random)
  environment_kwargs = environment_kwargs or {}
  return control.Environment(
      physics, task, time_limit=time_limit, control_timestep=_CONTROL_TIMESTEP,
      **environment_kwargs)


@SUITE.add('custom')
def urlb_flip(time_limit=_DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
  """Returns the Flip task."""
  physics = Physics.from_xml_string(*get_model_and_assets())
  task = Stickman(goal='urlb_flip', move_speed=_SPIN_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('custom')
def flipping(time_limit=_DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
  """Returns the Flipping task."""
  physics = Physics.from_xml_string(*get_model_and_assets())
  task = Stickman(goal='flipping', move_speed=2 * _RUN_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('custom')
def flip(time_limit=_DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
  """Returns the Flip task."""
  physics = Physics.from_xml_string(*get_model_and_assets())
  task = Stickman(goal='flip', move_speed=2 * _RUN_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('custom')
def backflip(time_limit=_DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
  """Returns the Backflip task."""
  physics = Physics.from_xml_string(*get_model_and_assets())
  task = Stickman(goal='flip', move_speed=-2 * _RUN_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('custom')
def stand(time_limit=_DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
  """Returns the Stand task."""
  physics = Physics.from_xml_string(*get_model_and_assets())
  task = Stickman(move_speed=0, goal='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('custom')
def walk(time_limit=_DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
  """Returns the Walk task."""
  physics = Physics.from_xml_string(*get_model_and_assets())
  task = Stickman(move_speed=_WALK_SPEED, goal='walk', random=random)
  environment_kwargs = environment_kwargs or {}
  return control.Environment(
      physics, task, time_limit=time_limit, control_timestep=_CONTROL_TIMESTEP,
      **environment_kwargs)


@SUITE.add('custom')
def run(time_limit=_DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
  """Returns the Run task."""
  physics = Physics.from_xml_string(*get_model_and_assets())
  task = Stickman(move_speed=_RUN_SPEED, goal='run', 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 stickman domain."""
    def torso_upright(self):
        """Returns projection from z-axes of torso to the z-axes of world."""
        return self.named.data.xmat['torso', 'zz']

    def torso_height(self):
        """Returns the height of the torso."""
        return self.named.data.xpos['torso', 'z']

    def horizontal_velocity(self):
        """Returns the horizontal velocity of the center-of-mass."""
        return self.named.data.sensordata['torso_subtreelinvel'][0]

    def orientations(self):
        """Returns planar orientations of all bodies."""
        return self.named.data.xmat[1:, ['xx', 'xz']].ravel()

    def angmomentum(self):
        """Returns the angular momentum of torso of the stickman about Y axis."""
        return self.named.data.subtree_angmom['torso'][1]


class Stickman(base.Task):
    """A planar stickman task."""
    def __init__(self, move_speed=0., goal='walk', forward=True, random=None):
        """Initializes an instance of `Stickman`.

    Args:
      move_speed: A float. If this value is zero, reward is given simply for
        standing up. Otherwise this specifies a target horizontal velocity for
        the walking task.
      random: Optional, either a `numpy.random.RandomState` instance, an
        integer seed for creating a new `RandomState`, or None to select a seed
        automatically (default).
    """
        self._move_speed = move_speed
        self._forward = 1 if forward else -1
        self._goal = goal
        super().__init__(random=random)

    def _hands_up_reward(self, physics):
        standing = self._stand_reward(physics)
        left_hand_height = physics.named.data.xpos['left_hand', 'z']
        right_hand_height = physics.named.data.xpos['right_hand', 'z']
        
        hand_height = (left_hand_height + right_hand_height) / 2

        hands_up = rewards.tolerance(hand_height,
                                bounds=(_YOGA_HANDS_UP_HEIGHT, float('inf')),
                                margin=_YOGA_HANDS_UP_HEIGHT/2)
        return standing * hands_up

    def _boxing_reward(self, physics):
        # torso up, but lower than standing
        # foot up, higher than torso
        # foot down
        standing = self._stand_reward(physics)

        left_hand_velocity = abs(physics.named.data.subtree_linvel['left_hand'][0])
        right_hand_velocity = abs(physics.named.data.subtree_linvel['right_hand'][0])
        punch_reward = rewards.tolerance(
                max(left_hand_velocity, right_hand_velocity),
                bounds=(_PUNCH_SPEED, float('inf')),
                margin=_PUNCH_SPEED / 2,
                value_at_margin=0.5,
                sigmoid='linear')

        # left_hand_dist = physics.named.data.xpos['left_hand', 'x'] - physics.named.data.xpos['torso', 'x']
        # right_hand_dist = physics.named.data.xpos['right_hand', 'x'] - physics.named.data.xpos['torso', 'x']
        # punch_reward = rewards.tolerance(
        #         max(left_hand_dist, right_hand_dist),
        #         bounds=(_PUNCH_DIST, float('inf')),
        #         margin=_PUNCH_DIST / 2,)

        return standing * punch_reward

    def _arabesque_reward(self, physics):
        # standing horizontal
        # one foot up, same height as torso
        # one foot down
        standing = rewards.tolerance(physics.torso_height(),
                                bounds=(_YOGA_STAND_HEIGHT, float('inf')),
                                margin=_YOGA_STAND_HEIGHT/2)
        
        left_foot_height = physics.named.data.xpos['left_foot', 'z']
        right_foot_height = physics.named.data.xpos['right_foot', 'z']
        
        max_foot = 'right_foot' if right_foot_height > left_foot_height else 'left_foot'
        min_foot = 'right_foot' if right_foot_height <= left_foot_height else 'left_foot'

        min_foot_height = physics.named.data.xpos[min_foot, 'z']
        max_foot_height = physics.named.data.xpos[max_foot, 'z']

        min_foot_down = rewards.tolerance(min_foot_height,
                                bounds=(-float('inf'), _YOGA_LIE_DOWN_HEIGHT),
                                margin=_YOGA_LIE_DOWN_HEIGHT*1.5)
        max_foot_up = rewards.tolerance(max_foot_height,
                                bounds=(_YOGA_STAND_HEIGHT, float('inf')),
                                margin=_YOGA_STAND_HEIGHT/2)
        
        min_foot_x = physics.named.data.xpos[min_foot, 'x']
        max_foot_x = physics.named.data.xpos[max_foot, 'x']
        
        correct_foot_pose = 0.1 if max_foot_x > min_foot_x else 1.0
 
        feet_pose = (min_foot_down + max_foot_up * 2) / 3
        return standing * feet_pose * correct_foot_pose
    
    def _lying_down_reward(self, physics):
        # torso down and horizontal
        # thigh and feet down
        torso_down = rewards.tolerance(physics.torso_height(),
                                bounds=(-float('inf'), _YOGA_LIE_DOWN_HEIGHT),
                                margin=_YOGA_LIE_DOWN_HEIGHT*1.5)
        horizontal = 1 - abs(physics.torso_upright())
        
        thigh_height = (physics.named.data.xpos['left_thigh', 'z'] + physics.named.data.xpos['right_thigh', 'z']) / 2
        thigh_down = rewards.tolerance(thigh_height,
                                bounds=(-float('inf'), _YOGA_LIE_DOWN_HEIGHT),
                                margin=_YOGA_LIE_DOWN_HEIGHT*1.5)
        leg_height = (physics.named.data.xpos['left_leg', 'z'] + physics.named.data.xpos['right_leg', 'z']) / 2
        leg_down = rewards.tolerance(leg_height,
                                bounds=(-float('inf'), _YOGA_LIE_DOWN_HEIGHT),
                                margin=_YOGA_LIE_DOWN_HEIGHT*1.5)
        feet_height = (physics.named.data.xpos['left_foot', 'z'] + physics.named.data.xpos['right_foot', 'z']) / 2
        feet_down = rewards.tolerance(feet_height,
                                bounds=(-float('inf'), _YOGA_LIE_DOWN_HEIGHT),
                                margin=_YOGA_LIE_DOWN_HEIGHT*1.5)
        return (3*torso_down + horizontal + thigh_down + feet_down + leg_down) / 7
    
    def _legs_up_reward(self, physics):
        # torso down and horizontal
        # legs up with thigh down
        torso_down = rewards.tolerance(physics.torso_height(),
                                bounds=(-float('inf'), _YOGA_LIE_DOWN_HEIGHT),
                                margin=_YOGA_LIE_DOWN_HEIGHT*1.5)
        horizontal = 1 - abs(physics.torso_upright())
        torso_down = (3*torso_down +horizontal) / 4
        
        feet_height = (physics.named.data.xpos['left_foot', 'z'] + physics.named.data.xpos['right_foot', 'z']) / 2
        feet_up = rewards.tolerance(feet_height,
                                bounds=(_YOGA_FEET_UP_LIE_DOWN_HEIGHT, float('inf')),
                                margin=_YOGA_FEET_UP_LIE_DOWN_HEIGHT/2)

        return torso_down * feet_up
    
    def _high_kick_reward(self, physics):
        # torso up, but lower than standing
        # foot up, higher than torso
        # foot down
        standing = rewards.tolerance(physics.torso_height(),
                                    bounds=(_YOGA_STAND_HEIGHT, float('inf')),
                                    margin=_YOGA_STAND_HEIGHT/2)

        left_foot_height = physics.named.data.xpos['left_foot', 'z']
        right_foot_height = physics.named.data.xpos['right_foot', 'z']
        
        min_foot_height = min(left_foot_height, right_foot_height)
        max_foot_height = max(left_foot_height, right_foot_height)

        min_foot_down = rewards.tolerance(min_foot_height,
                                bounds=(-float('inf'), _YOGA_LIE_DOWN_HEIGHT),
                                margin=_YOGA_LIE_DOWN_HEIGHT*1.5)
        max_foot_up = rewards.tolerance(max_foot_height,
                                bounds=(_STAND_HEIGHT, float('inf')),
                                margin=_STAND_HEIGHT/2)
        
        feet_pose = (3 * max_foot_up + min_foot_down) / 4

        return standing * feet_pose
    
    def _one_foot_reward(self, physics):
        # torso up, standing
        # foot up higher than foot down
        standing = rewards.tolerance(physics.torso_height(),
                                    bounds=(_YOGA_STAND_HEIGHT, float('inf')),
                                    margin=_YOGA_STAND_HEIGHT/2)

        left_foot_height = physics.named.data.xpos['left_foot', 'z']
        right_foot_height = physics.named.data.xpos['right_foot', 'z']
        
        min_foot_height = min(left_foot_height, right_foot_height)
        max_foot_height = max(left_foot_height, right_foot_height)

        min_foot_down = rewards.tolerance(min_foot_height,
                                bounds=(-float('inf'), _YOGA_LIE_DOWN_HEIGHT),
                                margin=_YOGA_LIE_DOWN_HEIGHT*1.5)
        max_foot_up = rewards.tolerance(max_foot_height,
                                bounds=(_YOGA_FEET_UP_HEIGHT, float('inf')),
                                margin=_YOGA_FEET_UP_HEIGHT/2)
        
        return standing * max_foot_up * min_foot_down

    def _lunge_pose_reward(self, physics):
        # torso up, standing, but lower
        # leg up higher than leg down
        # horiontal thigh and leg
        standing = rewards.tolerance(physics.torso_height(),
                                    bounds=(_YOGA_KNEESTAND_HEIGHT, float('inf')),
                                    margin=_YOGA_KNEESTAND_HEIGHT/2)
        upright = (1 + physics.torso_upright()) / 2
        torso = (3*standing + upright) / 4

        left_leg_height = physics.named.data.xpos['left_leg', 'z']
        right_leg_height = physics.named.data.xpos['right_leg', 'z']
        
        min_leg_height = min(left_leg_height, right_leg_height)
        max_leg_height = max(left_leg_height, right_leg_height)

        min_leg_down = rewards.tolerance(min_leg_height,
                                bounds=(-float('inf'), _YOGA_LIE_DOWN_HEIGHT),
                                margin=_YOGA_LIE_DOWN_HEIGHT*1.5)
        max_leg_up = rewards.tolerance(max_leg_height,
                                bounds=(_YOGA_KNEE_HEIGHT, float('inf')),
                                margin=_YOGA_KNEE_HEIGHT / 2)
        
        max_thigh = 'left_thigh' if max_leg_height == left_leg_height else 'right_thigh'
        min_leg = 'left_leg' if min_leg_height == left_leg_height else 'right_leg'

        max_thigh_horiz = 1 - abs(physics.named.data.xmat[max_thigh, 'zz'])
        min_leg_horiz = 1 - abs(physics.named.data.xmat[min_leg, 'zz'])
        
        legs = (min_leg_down + max_leg_up + max_thigh_horiz + min_leg_horiz) / 4

        return torso * legs

    def _sit_knees_reward(self, physics):
        # torso up, standing, but lower
        # foot up higher than foot down
        standing = rewards.tolerance(physics.torso_height(),
                                    bounds=(_YOGA_SITTING_HEIGHT, float('inf')),
                                    margin=_YOGA_SITTING_HEIGHT/2)
        upright = (1 + physics.torso_upright()) / 2
        torso_up = (3*standing + upright) / 4

        legs_height = (physics.named.data.xpos['left_leg', 'z'] + physics.named.data.xpos['right_leg', 'z']) / 2
        legs_down = rewards.tolerance(legs_height,
                                bounds=(-float('inf'), _YOGA_SITTING_LEGS_HEIGHT),
                                margin=_YOGA_SITTING_LEGS_HEIGHT*1.5)
        feet_height = (physics.named.data.xpos['left_foot', 'z'] + physics.named.data.xpos['right_foot', 'z']) / 2
        feet_down = rewards.tolerance(feet_height,
                                bounds=(-float('inf'), _YOGA_LIE_DOWN_HEIGHT),
                                margin=_YOGA_LIE_DOWN_HEIGHT*1.5)
        
        l_thigh_foot_distance = max(0.1, abs(physics.named.data.xpos['left_foot', 'x'] - physics.named.data.xpos['left_thigh', 'x'])) - 0.1
        r_thigh_foot_distance = max(0.1, abs(physics.named.data.xpos['right_foot', 'x'] - physics.named.data.xpos['right_thigh', 'x'])) - 0.1
        close = np.exp(-(l_thigh_foot_distance + r_thigh_foot_distance)/2)
        
        legs = (3 * legs_down + feet_down) / 4
        return torso_up * legs * close

    def _urlb_flip_reward(self, physics):
        standing = rewards.tolerance(physics.torso_height(),
                                     bounds=(_STAND_HEIGHT, float('inf')),
                                     margin=_STAND_HEIGHT / 2)
        upright = (1 + physics.torso_upright()) / 2
        stand_reward = (3 * standing + upright) / 4
        move_reward = rewards.tolerance(self._forward *
                                        physics.named.data.subtree_angmom['torso'][1], # physics.angmomentum(),
                                        bounds=(_SPIN_SPEED, float('inf')),
                                        margin=_SPIN_SPEED,
                                        value_at_margin=0,
                                        sigmoid='linear')
        return stand_reward * (5 * move_reward + 1) / 6

    def _flip_reward(self, physics):
        thigh_height = (physics.named.data.xpos['left_thigh', 'z'] + physics.named.data.xpos['right_thigh', 'z']) / 2
        thigh_up = rewards.tolerance(thigh_height,
                                bounds=(_YOGA_STAND_HEIGHT, float('inf')),
                                margin=_YOGA_STAND_HEIGHT/2)
        feet_height = (physics.named.data.xpos['left_foot', 'z'] + physics.named.data.xpos['right_foot', 'z']) / 2
        legs_up = rewards.tolerance(feet_height,
                                bounds=(_YOGA_LEGS_UP_HEIGHT, float('inf')),
                                margin=_YOGA_LEGS_UP_HEIGHT/2)
        upside_down_reward = (3*legs_up + 2*thigh_up) / 5
        if self._move_speed == 0:
            return upside_down_reward
        move_reward = rewards.tolerance(physics.named.data.subtree_angmom['torso'][1], # physics.angmomentum(),
                                    bounds=(self._move_speed, float('inf')) if self._move_speed > 0 else (-float('inf'), self._move_speed),
                                    margin=abs(self._move_speed)/2,
                                    value_at_margin=0.5,
                                    sigmoid='linear')
        return upside_down_reward * (5*move_reward + 1) / 6

    
    def _stand_reward(self, physics):
        standing = rewards.tolerance(physics.torso_height(),
                                    bounds=(_STAND_HEIGHT, float('inf')),
                                    margin=_STAND_HEIGHT / 2)
        upright = (1 + physics.torso_upright()) / 2
        return (3 * standing + upright) / 4

    def initialize_episode(self, physics):
        """Sets the state of the environment at the start of each episode.

    In 'standing' mode, use initial orientation and small velocities.
    In 'random' mode, randomize joint angles and let fall to the floor.

    Args:
      physics: An instance of `Physics`.

    """
        randomizers.randomize_limited_and_rotational_joints(physics, self.random)
        super().initialize_episode(physics)

    def get_observation(self, physics):
        """Returns an observation of body orientations, height and velocites."""
        obs = collections.OrderedDict()
        obs['orientations'] = physics.orientations()
        obs['height'] = physics.torso_height()
        obs['velocity'] = physics.velocity()
        return obs

    def get_reward(self, physics):
        """Returns a reward to the agent."""
        if self._goal in ['stand', 'walk', 'run']:
            stand_reward = self._stand_reward(physics)
            move_reward = rewards.tolerance(
                self._forward * physics.horizontal_velocity(),
                bounds=(self._move_speed, float('inf')),
                margin=self._move_speed / 2,
                value_at_margin=0.5,
                sigmoid='linear')
            return stand_reward * (5 * move_reward + 1) / 6
        if self._goal == 'flipping':
            self._move_speed = abs(self._move_speed)
            pos_rew = self._flip_reward(physics)
            self._move_speed = -abs(self._move_speed)
            neg_rew = self._flip_reward(physics)
            return max(pos_rew, neg_rew)
        try:
            reward_fn = getattr(self, f'_{self._goal}_reward')
            return reward_fn(physics)
        except Exception as e:
            print(e)
            raise NotImplementedError(f'Goal {self._goal} or function "_{self._goal}_reward" not implemented.')

if __name__ == '__main__':
    from dm_control import viewer
    import numpy as np
    
    env = boxing()
    env.task.visualize_reward = True

    action_spec = env.action_spec()

    def zero_policy(time_step):
        print(time_step.reward)
        return np.zeros(action_spec.shape)
    
    ts = env.reset()
    while True:
       ts = env.step(zero_policy(ts))

    viewer.launch(env, policy=zero_policy)
    
    # obs = env.reset()
    # next_obs, reward, done, info = env.step(np.zeros(6))