# Copyright 2019 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.

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

"""Tasks where the goal is to elevate a prop."""

import collections

import itertools

from dm_control import composer

from dm_control.composer import initializers

from dm_control.composer.variation import distributions

from dm_control.entities import props

from dm_control.manipulation.shared import arenas

from dm_control.manipulation.shared import cameras

from dm_control.manipulation.shared import constants

from dm_control.manipulation.shared import observations

from dm_control.manipulation.shared import registry

from dm_control.manipulation.shared import robots

from dm_control.manipulation.shared import tags

from dm_control.manipulation.shared import workspaces

from dm_control.utils import rewards

import numpy as np

_LiftWorkspace = collections.namedtuple(

'_LiftWorkspace', ['prop_bbox', 'tcp_bbox', 'arm_offset'])

_DUPLO_WORKSPACE = _LiftWorkspace(

prop_bbox=workspaces.BoundingBox(

lower=(-0.1, -0.1, 0.0),

upper=(0.1, 0.1, 0.0)),

tcp_bbox=workspaces.BoundingBox(

lower=(-0.1, -0.1, 0.2),

upper=(0.1, 0.1, 0.4)),

arm_offset=robots.ARM_OFFSET)

_BOX_SIZE = 0.09

_BOX_MASS = 1.3

_BOX_WORKSPACE = _LiftWorkspace(

prop_bbox=workspaces.BoundingBox(

lower=(-0.1, -0.1, _BOX_SIZE),

upper=(0.1, 0.1, _BOX_SIZE)),

tcp_bbox=workspaces.BoundingBox(

lower=(-0.1, -0.1, 0.2),

upper=(0.1, 0.1, 0.4)),

arm_offset=robots.ARM_OFFSET)

_DISTANCE_TO_LIFT = 0.3

class _VertexSitesMixin:

"""Mixin class that adds sites corresponding to the vertices of a box."""

def _add_vertex_sites(self, box_geom_or_site):

"""Add sites corresponding to the vertices of a box geom or site."""

offsets = (

(-half_length, half_length) for half_length in box_geom_or_site.size)

site_positions = np.vstack(list(itertools.product(*offsets)))

if box_geom_or_site.pos is not None:

site_positions += box_geom_or_site.pos

self._vertices = []

for i, pos in enumerate(site_positions):

site = box_geom_or_site.parent.add(

'site', name='vertex_' + str(i), pos=pos, type='sphere', size=[0.002],

rgba=constants.RED, group=constants.TASK_SITE_GROUP)

self._vertices.append(site)

@property

def vertices(self):

return self._vertices

class _BoxWithVertexSites(props.Primitive, _VertexSitesMixin):

"""Subclass of `Box` with sites marking the vertices of the box geom."""

def _build(self, *args, **kwargs):

super()._build(*args, geom_type='box', **kwargs)

self._add_vertex_sites(self.geom)

class _DuploWithVertexSites(props.Duplo, _VertexSitesMixin):

"""Subclass of `Duplo` with sites marking the vertices of its sensor site."""

def _build(self, *args, **kwargs):

super()._build(*args, **kwargs)

self._add_vertex_sites(self.mjcf_model.find('site', 'bounding_box'))

class Lift(composer.Task):

"""A task where the goal is to elevate a prop."""

def __init__(

self, arena, arm, hand, prop, obs_settings, workspace, control_timestep):

"""Initializes a new `Lift` task.

Args:

arena: `composer.Entity` instance.

arm: `robot_base.RobotArm` instance.

hand: `robot_base.RobotHand` instance.

prop: `composer.Entity` instance.

obs_settings: `observations.ObservationSettings` instance.

workspace: `_LiftWorkspace` specifying the placement of the prop and TCP.

control_timestep: Float specifying the control timestep in seconds.

"""

self._arena = arena

self._arm = arm

self._hand = hand

self._arm.attach(self._hand)

self._arena.attach_offset(self._arm, offset=workspace.arm_offset)

self.control_timestep = control_timestep

# Add custom camera observable.

self._task_observables = cameras.add_camera_observables(

arena, obs_settings, cameras.FRONT_CLOSE)

self._tcp_initializer = initializers.ToolCenterPointInitializer(

self._hand, self._arm,

position=distributions.Uniform(*workspace.tcp_bbox),

quaternion=workspaces.DOWN_QUATERNION)

self._prop = prop

self._arena.add_free_entity(prop)

self._prop_placer = initializers.PropPlacer(

props=[prop],

position=distributions.Uniform(*workspace.prop_bbox),

quaternion=workspaces.uniform_z_rotation,

ignore_collisions=True,

settle_physics=True)

# Add sites for visualizing bounding boxes and target height.

self._target_height_site = workspaces.add_bbox_site(

body=self.root_entity.mjcf_model.worldbody,

lower=(-1, -1, 0), upper=(1, 1, 0),

rgba=constants.RED, name='target_height')

workspaces.add_bbox_site(

body=self.root_entity.mjcf_model.worldbody,

lower=workspace.tcp_bbox.lower, upper=workspace.tcp_bbox.upper,

rgba=constants.GREEN, name='tcp_spawn_area')

workspaces.add_bbox_site(

body=self.root_entity.mjcf_model.worldbody,

lower=workspace.prop_bbox.lower, upper=workspace.prop_bbox.upper,

rgba=constants.BLUE, name='prop_spawn_area')

@property

def root_entity(self):

return self._arena

@property

def arm(self):

return self._arm

@property

def hand(self):

return self._hand

@property

def task_observables(self):

return self._task_observables

def _get_height_of_lowest_vertex(self, physics):

return min(physics.bind(self._prop.vertices).xpos[:, 2])

def get_reward(self, physics):

prop_height = self._get_height_of_lowest_vertex(physics)

return rewards.tolerance(prop_height,

bounds=(self._target_height, np.inf),

margin=_DISTANCE_TO_LIFT,

value_at_margin=0,

sigmoid='linear')

def initialize_episode(self, physics, random_state):

self._hand.set_grasp(physics, close_factors=random_state.uniform())

self._prop_placer(physics, random_state)

self._tcp_initializer(physics, random_state)

# Compute the target height based on the initial height of the prop's

# center of mass after settling.

initial_prop_height = self._get_height_of_lowest_vertex(physics)

self._target_height = _DISTANCE_TO_LIFT + initial_prop_height

physics.bind(self._target_height_site).pos[2] = self._target_height

def _lift(obs_settings, prop_name):

"""Configure and instantiate a Lift task.

Args:

obs_settings: `observations.ObservationSettings` instance.

prop_name: The name of the prop to be lifted. Must be either 'duplo' or

'box'.

Returns:

An instance of `lift.Lift`.

Raises:

ValueError: If `prop_name` is neither 'duplo' nor 'box'.

"""

arena = arenas.Standard()

arm = robots.make_arm(obs_settings=obs_settings)

hand = robots.make_hand(obs_settings=obs_settings)

if prop_name == 'duplo':

workspace = _DUPLO_WORKSPACE

prop = _DuploWithVertexSites(

observable_options=observations.make_options(

obs_settings, observations.FREEPROP_OBSERVABLES))

elif prop_name == 'box':

workspace = _BOX_WORKSPACE

# NB: The box is intentionally too large to pick up with a pinch grip.

prop = _BoxWithVertexSites(

size=[_BOX_SIZE] * 3,

observable_options=observations.make_options(

obs_settings, observations.FREEPROP_OBSERVABLES))

prop.geom.mass = _BOX_MASS

else:

raise ValueError('`prop_name` must be either \'duplo\' or \'box\'.')

task = Lift(arena=arena, arm=arm, hand=hand, prop=prop, workspace=workspace,

obs_settings=obs_settings,

control_timestep=constants.CONTROL_TIMESTEP)

return task

@registry.add(tags.FEATURES)

def lift_brick_features():

return _lift(obs_settings=observations.PERFECT_FEATURES, prop_name='duplo')

@registry.add(tags.VISION)

def lift_brick_vision():

return _lift(obs_settings=observations.VISION, prop_name='duplo')

@registry.add(tags.FEATURES)

def lift_large_box_features():

return _lift(obs_settings=observations.PERFECT_FEATURES, prop_name='box')

@registry.add(tags.VISION)

def lift_large_box_vision():

return _lift(obs_settings=observations.VISION, prop_name='box')