Spaces:
Running on Zero
Running on Zero
File size: 2,926 Bytes
23a59ea | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 | import os
from dm_control.rl import control
from dm_control.suite import base
from dm_control.suite import common
from dm_control.suite import fish
from dm_control.utils import rewards
from dm_control.utils import io as resources
import numpy as np
_TASKS_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'dmcontrol')
_DEFAULT_TIME_LIMIT = 40
_CONTROL_TIMESTEP = .04
_JOINTS = ['tail1',
'tail_twist',
'tail2',
'finright_roll',
'finright_pitch',
'finleft_roll',
'finleft_pitch']
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, 'fish.xml')), common.ASSETS
@fish.SUITE.add('custom')
def obstacles(time_limit=_DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
"""Returns the Fish Obstacles task."""
physics = fish.Physics.from_xml_string(*get_model_and_assets())
task = Obstacles(random=random)
environment_kwargs = environment_kwargs or {}
return control.Environment(
physics, task, control_timestep=_CONTROL_TIMESTEP, time_limit=time_limit,
**environment_kwargs)
class Obstacles(fish.Swim):
"""A custom Fish Obstacles task."""
def __init__(self, random=None):
super().__init__(random=random)
def in_wall(self, physics, name, min_distance=0.08):
"""Returns True if the given body is too close to a wall."""
for wall in ['wall0', 'wall1', 'wall2', 'wall3']:
l1_dist = np.min(np.abs(physics.named.data.geom_xpos[name][:2] - physics.named.data.geom_xpos[wall][:2]))
if l1_dist < min_distance:
return True
return False
def initialize_episode(self, physics):
in_wall = True
while in_wall:
# Randomize fish position.
quat = self.random.randn(4)
physics.named.data.qpos['root'][3:7] = quat / np.linalg.norm(quat)
for joint in _JOINTS:
physics.named.data.qpos[joint] = self.random.uniform(-.2, .2)
# Randomize target position.
physics.named.model.geom_pos['target', 'x'] = self.random.uniform(-.4, .4)
physics.named.model.geom_pos['target', 'y'] = self.random.uniform(-.4, .4)
physics.named.model.geom_pos['target', 'z'] = self.random.uniform(.1, .3)
# Make sure target is not too close to a wall.
physics.after_reset()
in_wall = self.in_wall(physics, 'target')
base.Task.initialize_episode(self, physics)
def get_reward(self, physics):
radii = physics.named.model.geom_size[['mouth', 'target'], 0].sum()
in_target = rewards.tolerance(np.linalg.norm(physics.mouth_to_target()),
bounds=(0, radii), margin=2*radii)
is_upright = 0.5 * (physics.upright() + 1)
is_not_in_wall = 1. - self.in_wall(physics, 'torso', min_distance=0.06)
return is_not_in_wall * (7*in_target + is_upright) / 8
|