envs/custom_dmc_tasks/walker.py

import os

import numpy as np
from dm_control.rl import control
from dm_control.suite import common
from dm_control.suite import walker
from dm_control.utils import rewards
from dm_control.utils import io as resources

_TASKS_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'custom_dmc_tasks')

_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/walker.py
_SPIN_SPEED = 5.0
#

class WalkerYogaPoses:
    """
    Joint positions for some yoga poses
    """
    lie_back   = [ -1.2 ,  0. ,  -1.57,  0, 0. , 0.0, 0, -0.,  0.0]
    lie_front  = [-1.2,   -0,      1.57, 0, -0.2, 0, 0, -0.2, 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_foot   = [-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 ]


def get_model_and_assets():
    """Returns a tuple containing the model XML string and a dict of assets."""
    return resources.GetResource(os.path.join(_TASKS_DIR, 'walker.xml')), common.ASSETS


@walker.SUITE.add('custom')
def walk_backwards(time_limit=walker._DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
  """Returns the Walk Backwards task."""
  physics = walker.Physics.from_xml_string(*get_model_and_assets())
  task = BackwardsPlanarWalker(move_speed=walker._WALK_SPEED, random=random)
  environment_kwargs = environment_kwargs or {}
  return control.Environment(
      physics, task, time_limit=time_limit, control_timestep=walker._CONTROL_TIMESTEP,
      **environment_kwargs)


@walker.SUITE.add('custom')
def run_backwards(time_limit=walker._DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
  """Returns the Run Backwards task."""
  physics = walker.Physics.from_xml_string(*get_model_and_assets())
  task = BackwardsPlanarWalker(move_speed=walker._RUN_SPEED, random=random)
  environment_kwargs = environment_kwargs or {}
  return control.Environment(
      physics, task, time_limit=time_limit, control_timestep=walker._CONTROL_TIMESTEP,
      **environment_kwargs)


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


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


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

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

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

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

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

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


@walker.SUITE.add('custom')
def urlb_flip(time_limit=walker._DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
  """Returns the Flip task."""
  physics = walker.Physics.from_xml_string(*get_model_and_assets())
  task = YogaPlanarWalker(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=walker._CONTROL_TIMESTEP,
      **environment_kwargs)


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

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


@walker.SUITE.add('custom')
def backflip(time_limit=walker._DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
  """Returns the Backflip task."""
  physics = walker.Physics.from_xml_string(*get_model_and_assets())
  task = YogaPlanarWalker(goal='flip', move_speed=-2 * walker._RUN_SPEED, random=random)
  environment_kwargs = environment_kwargs or {}
  return control.Environment(
      physics, task, time_limit=time_limit, control_timestep=walker._CONTROL_TIMESTEP,
      **environment_kwargs)


class BackwardsPlanarWalker(walker.PlanarWalker):
    """Backwards PlanarWalker task."""
    def __init__(self, move_speed, random=None):
        super().__init__(move_speed, random)
    
    def get_reward(self, physics):
        standing = rewards.tolerance(physics.torso_height(),
                                 bounds=(_YOGA_STAND_HEIGHT, float('inf')),
                                 margin=_YOGA_STAND_HEIGHT/2)
        upright = (1 + physics.torso_upright()) / 2
        stand_reward = (3*standing + upright) / 4
        if self._move_speed == 0:
            return stand_reward
        else:
            move_reward = rewards.tolerance(physics.horizontal_velocity(),
                                            bounds=(-float('inf'), -self._move_speed),
                                            margin=self._move_speed/2,
                                            value_at_margin=0.5,
                                            sigmoid='linear')
            return stand_reward * (5*move_reward + 1) / 6


class YogaPlanarWalker(walker.PlanarWalker):
    """Yoga PlanarWalker tasks."""
    
    def __init__(self, goal='arabesque', move_speed=0, random=None):
        super().__init__(0, random)
        self._goal = goal
        self._move_speed = move_speed
    
    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=(walker._STAND_HEIGHT, float('inf')),
                                margin=walker._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=(walker._STAND_HEIGHT, float('inf')),
                                     margin=walker._STAND_HEIGHT / 2)
        upright = (1 + physics.torso_upright()) / 2
        stand_reward = (3 * standing + upright) / 4
        move_reward = rewards.tolerance(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 get_reward(self, physics):
        if self._goal == 'arabesque':
            return self._arabesque_reward(physics)
        elif self._goal == 'lying_down':
            return self._lying_down_reward(physics)
        elif self._goal == 'legs_up':
            return self._legs_up_reward(physics)
        elif self._goal == 'flip':
            return self._flip_reward(physics)
        elif 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)
        elif self._goal == 'high_kick':
            return self._high_kick_reward(physics)
        elif self._goal == 'one_foot':
            return self._one_foot_reward(physics)
        elif self._goal == 'lunge_pose':
            return self._lunge_pose_reward(physics)
        elif self._goal == 'sit_knees':
            return self._sit_knees_reward(physics)
        elif self._goal == 'urlb_flip':
            return self._urlb_flip_reward(physics)
        else:
            raise NotImplementedError(f'Goal {self._goal} is not implemented.')


if __name__ == '__main__':
    from dm_control import viewer
    import numpy as np
    
    env = sit_knees()
    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)
    viewer.launch(env, policy=zero_policy)
    
    # obs = env.reset()
    # next_obs, reward, done, info = env.step(np.zeros(6))