Skip to content

Commit 3301c73

Browse files
jsmerelDMcopybara-github
jsmerelDM
authored and
copybara-github
committed
Initial open-sourcing of rodent and tasks.
PiperOrigin-RevId: 295778102 Change-Id: I25775d94c799989760350f09a89226196939c3b7
1 parent cec6981 commit 3301c73

19 files changed

Lines changed: 2166 additions & 7 deletions

‎dm_control/locomotion/README.md‎

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -12,9 +12,6 @@ available from this library.
1212
<img src="gaps.png" height="200">
1313
</p>
1414

15-
**NOTE:** Rodent walker model and examples will be made available as part of
16-
this library before the ICLR 2020 conference.
17-
1815
## Terminology
1916

2017
This library facilitates the creation of environments that require **walkers**
@@ -80,8 +77,7 @@ papers. Relevant references include:
8077

8178
- [Neural probabilistic motor primitives for humanoid control (2019)][merel2019b].
8279

83-
- [Deep neuroethology of a virtual rodent (2020)][merel2020] -- code coming
84-
soon.
80+
- [Deep neuroethology of a virtual rodent (2020)][merel2020].
8581

8682
[installation-and-requirements]: ../../README.md#installation-and-requirements
8783
[`dm_control.viewer`]: ../viewer/README.md

‎dm_control/locomotion/arenas/__init__.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,9 +14,13 @@
1414
# ============================================================================
1515
"""Arenas for Locomotion tasks."""
1616

17+
from dm_control.locomotion.arenas.bowl import Bowl
1718
from dm_control.locomotion.arenas.corridors import EmptyCorridor
1819
from dm_control.locomotion.arenas.corridors import GapsCorridor
1920
from dm_control.locomotion.arenas.corridors import WallsCorridor
2021
from dm_control.locomotion.arenas.floors import Floor
22+
from dm_control.locomotion.arenas.labmaze_textures import FloorTextures
23+
from dm_control.locomotion.arenas.labmaze_textures import SkyBox
24+
from dm_control.locomotion.arenas.labmaze_textures import WallTextures
2125
from dm_control.locomotion.arenas.mazes import MazeWithTargets
2226
from dm_control.locomotion.arenas.mazes import RandomMazeWithTargets

‎dm_control/locomotion/arenas/bowl.py‎

Lines changed: 139 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,139 @@
1+
# Copyright 2020 The dm_control Authors.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
# ============================================================================
15+
"""Bowl arena with bumps."""
16+
17+
from __future__ import absolute_import
18+
from __future__ import division
19+
from __future__ import print_function
20+
21+
from dm_control import composer
22+
from dm_control.locomotion.arenas import assets as locomotion_arenas_assets
23+
from dm_control.mujoco.wrapper import mjbindings
24+
25+
import numpy as np
26+
from scipy import ndimage
27+
28+
mjlib = mjbindings.mjlib
29+
30+
_TOP_CAMERA_DISTANCE = 100
31+
_TOP_CAMERA_Y_PADDING_FACTOR = 1.1
32+
33+
# Constants related to terrain generation.
34+
_TERRAIN_SMOOTHNESS = .5 # 0.0: maximally bumpy; 1.0: completely smooth.
35+
_TERRAIN_BUMP_SCALE = .2 # Spatial scale of terrain bumps (in meters).
36+
37+
38+
class Bowl(composer.Arena):
39+
"""A bowl arena with sinusoidal bumps."""
40+
41+
def _build(self, size=(10, 10), aesthetic='default', name='bowl'):
42+
super(Bowl, self)._build(name=name)
43+
44+
self._hfield = self._mjcf_root.asset.add(
45+
'hfield',
46+
name='terrain',
47+
nrow=201,
48+
ncol=201,
49+
size=(6, 6, 0.5, 0.1))
50+
51+
if aesthetic != 'default':
52+
ground_info = locomotion_arenas_assets.get_ground_texture_info(aesthetic)
53+
sky_info = locomotion_arenas_assets.get_sky_texture_info(aesthetic)
54+
texturedir = locomotion_arenas_assets.get_texturedir(aesthetic)
55+
self._mjcf_root.compiler.texturedir = texturedir
56+
57+
self._texture = self._mjcf_root.asset.add(
58+
'texture', name='aesthetic_texture', file=ground_info.file,
59+
type=ground_info.type)
60+
self._material = self._mjcf_root.asset.add(
61+
'material', name='aesthetic_material', texture=self._texture,
62+
texuniform='true')
63+
self._skybox = self._mjcf_root.asset.add(
64+
'texture', name='aesthetic_skybox', file=sky_info.file,
65+
type='skybox', gridsize=sky_info.gridsize,
66+
gridlayout=sky_info.gridlayout)
67+
self._terrain_geom = self._mjcf_root.worldbody.add(
68+
'geom',
69+
name='terrain',
70+
type='hfield',
71+
pos=(0, 0, -0.01),
72+
hfield='terrain',
73+
material=self._material)
74+
self._ground_geom = self._mjcf_root.worldbody.add(
75+
'geom',
76+
type='plane',
77+
name='groundplane',
78+
size=list(size) + [0.5],
79+
material=self._material)
80+
else:
81+
self._terrain_geom = self._mjcf_root.worldbody.add(
82+
'geom',
83+
name='terrain',
84+
type='hfield',
85+
rgba=(0.2, 0.3, 0.4, 1),
86+
pos=(0, 0, -0.01),
87+
hfield='terrain')
88+
self._ground_geom = self._mjcf_root.worldbody.add(
89+
'geom',
90+
type='plane',
91+
name='groundplane',
92+
rgba=(0.2, 0.3, 0.4, 1),
93+
size=list(size) + [0.5])
94+
95+
self._mjcf_root.visual.headlight.set_attributes(
96+
ambient=[.4, .4, .4], diffuse=[.8, .8, .8], specular=[.1, .1, .1])
97+
98+
self._regenerate = True
99+
100+
def regenerate(self, random_state):
101+
# regeneration of the bowl requires physics, so postponed to initialization.
102+
self._regenerate = True
103+
104+
def initialize_episode(self, physics, random_state):
105+
if self._regenerate:
106+
self._regenerate = False
107+
108+
# Get heightfield resolution, assert that it is square.
109+
res = physics.bind(self._hfield).nrow
110+
assert res == physics.bind(self._hfield).ncol
111+
112+
# Sinusoidal bowl shape.
113+
row_grid, col_grid = np.ogrid[-1:1:res*1j, -1:1:res*1j]
114+
radius = np.clip(np.sqrt(col_grid**2 + row_grid**2), .1, 1)
115+
bowl_shape = .5 - np.cos(2*np.pi*radius)/2
116+
117+
# Random smooth bumps.
118+
terrain_size = 2 * physics.bind(self._hfield).size[0]
119+
bump_res = int(terrain_size / _TERRAIN_BUMP_SCALE)
120+
bumps = random_state.uniform(_TERRAIN_SMOOTHNESS, 1, (bump_res, bump_res))
121+
smooth_bumps = ndimage.zoom(bumps, res / float(bump_res))
122+
123+
# Terrain is elementwise product.
124+
terrain = bowl_shape * smooth_bumps
125+
start_idx = physics.bind(self._hfield).adr
126+
physics.model.hfield_data[start_idx:start_idx+res**2] = terrain.ravel()
127+
128+
# If we have a rendering context, we need to re-upload the modified
129+
# heightfield data.
130+
if physics.contexts:
131+
with physics.contexts.gl.make_current() as ctx:
132+
ctx.call(mjlib.mjr_uploadHField,
133+
physics.model.ptr,
134+
physics.contexts.mujoco.ptr,
135+
physics.bind(self._hfield).element_id)
136+
137+
@property
138+
def ground_geoms(self):
139+
return (self._terrain_geom, self._ground_geom)

‎dm_control/locomotion/arenas/bowl_test.py‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
# Copyright 2020 The dm_control Authors.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
# ============================================================================
15+
"""Tests for locomotion.arenas.bowl."""
16+
17+
from __future__ import absolute_import
18+
from __future__ import division
19+
from __future__ import print_function
20+
21+
from absl.testing import absltest
22+
from dm_control import mjcf
23+
from dm_control.locomotion.arenas import bowl
24+
25+
26+
class BowlTest(absltest.TestCase):
27+
28+
def test_can_compile_mjcf(self):
29+
30+
arena = bowl.Bowl()
31+
mjcf.Physics.from_mjcf_model(arena.mjcf_model)
32+
33+
34+
if __name__ == '__main__':
35+
absltest.main()

‎dm_control/locomotion/examples/basic_rodent_2020.py‎

Lines changed: 173 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,173 @@
1+
# Copyright 2020 The dm_control Authors.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
# ============================================================================
15+
"""Produces reference environments for rodent tasks."""
16+
17+
from __future__ import absolute_import
18+
from __future__ import division
19+
from __future__ import print_function
20+
21+
import functools
22+
23+
from dm_control import composer
24+
from dm_control.composer.variation import distributions
25+
from dm_control.locomotion.arenas import bowl
26+
from dm_control.locomotion.arenas import corridors as corr_arenas
27+
from dm_control.locomotion.arenas import floors
28+
from dm_control.locomotion.arenas import labmaze_textures
29+
from dm_control.locomotion.arenas import mazes
30+
from dm_control.locomotion.props import target_sphere
31+
from dm_control.locomotion.tasks import corridors as corr_tasks
32+
from dm_control.locomotion.tasks import escape
33+
from dm_control.locomotion.tasks import random_goal_maze
34+
from dm_control.locomotion.tasks import reach
35+
from dm_control.locomotion.walkers import rodent
36+
37+
_CONTROL_TIMESTEP = .02
38+
_PHYSICS_TIMESTEP = 0.001
39+
40+
41+
def rodent_escape_bowl(random_state=None):
42+
"""Requires a rodent to climb out of a bowl-shaped terrain."""
43+
44+
# Build a position-controlled rodent walker.
45+
walker = rodent.Rat(
46+
observable_options={'egocentric_camera': dict(enabled=True)})
47+
48+
# Build a corridor-shaped arena that is obstructed by walls.
49+
arena = bowl.Bowl(
50+
size=(20., 20.),
51+
aesthetic='outdoor_natural')
52+
53+
# Build a task that rewards the agent for running down the corridor at a
54+
# specific velocity.
55+
task = escape.Escape(
56+
walker=walker,
57+
arena=arena,
58+
physics_timestep=_PHYSICS_TIMESTEP,
59+
control_timestep=_CONTROL_TIMESTEP)
60+
61+
return composer.Environment(time_limit=20,
62+
task=task,
63+
random_state=random_state,
64+
strip_singleton_obs_buffer_dim=True)
65+
66+
67+
def rodent_run_gaps(random_state=None):
68+
"""Requires a rodent to run down a corridor with gaps."""
69+
70+
# Build a position-controlled rodent walker.
71+
walker = rodent.Rat(
72+
observable_options={'egocentric_camera': dict(enabled=True)})
73+
74+
# Build a corridor-shaped arena with gaps, where the sizes of the gaps and
75+
# platforms are uniformly randomized.
76+
arena = corr_arenas.GapsCorridor(
77+
platform_length=distributions.Uniform(.4, .8),
78+
gap_length=distributions.Uniform(.05, .2),
79+
corridor_width=2,
80+
corridor_length=40,
81+
aesthetic='outdoor_natural')
82+
83+
# Build a task that rewards the agent for running down the corridor at a
84+
# specific velocity.
85+
task = corr_tasks.RunThroughCorridor(
86+
walker=walker,
87+
arena=arena,
88+
walker_spawn_position=(5, 0, 0),
89+
walker_spawn_rotation=0,
90+
target_velocity=1.0,
91+
contact_termination=False,
92+
terminate_at_height=-0.3,
93+
physics_timestep=_PHYSICS_TIMESTEP,
94+
control_timestep=_CONTROL_TIMESTEP)
95+
96+
return composer.Environment(time_limit=30,
97+
task=task,
98+
random_state=random_state,
99+
strip_singleton_obs_buffer_dim=True)
100+
101+
102+
def rodent_maze_forage(random_state=None):
103+
"""Requires a rodent to find all items in a maze."""
104+
105+
# Build a position-controlled rodent walker.
106+
walker = rodent.Rat(
107+
observable_options={'egocentric_camera': dict(enabled=True)})
108+
109+
# Build a maze with rooms and targets.
110+
wall_textures = labmaze_textures.WallTextures(style='style_01')
111+
arena = mazes.RandomMazeWithTargets(
112+
x_cells=11,
113+
y_cells=11,
114+
xy_scale=.5,
115+
z_height=.3,
116+
max_rooms=4,
117+
room_min_size=4,
118+
room_max_size=5,
119+
spawns_per_room=1,
120+
targets_per_room=3,
121+
wall_textures=wall_textures,
122+
aesthetic='outdoor_natural')
123+
124+
# Build a task that rewards the agent for obtaining targets.
125+
task = random_goal_maze.ManyGoalsMaze(
126+
walker=walker,
127+
maze_arena=arena,
128+
target_builder=functools.partial(
129+
target_sphere.TargetSphere,
130+
radius=0.05,
131+
height_above_ground=.125,
132+
rgb1=(0, 0, 0.4),
133+
rgb2=(0, 0, 0.7)),
134+
target_reward_scale=50.,
135+
contact_termination=False,
136+
physics_timestep=_PHYSICS_TIMESTEP,
137+
control_timestep=_CONTROL_TIMESTEP)
138+
139+
return composer.Environment(time_limit=30,
140+
task=task,
141+
random_state=random_state,
142+
strip_singleton_obs_buffer_dim=True)
143+
144+
145+
def rodent_two_touch(random_state=None):
146+
"""Requires a rodent to tap an orb, wait an interval, and tap it again."""
147+
148+
# Build a position-controlled rodent walker.
149+
walker = rodent.Rat(
150+
observable_options={'egocentric_camera': dict(enabled=True)})
151+
152+
arena = floors.Floor(
153+
size=(10., 10.),
154+
aesthetic='outdoor_natural')
155+
156+
task = reach.TwoTouch(
157+
walker=walker,
158+
arena=arena,
159+
target_builders=[
160+
functools.partial(target_sphere.TargetSphereTwoTouch, radius=0.025),
161+
],
162+
randomize_spawn_rotation=True,
163+
target_type_rewards=[25.],
164+
shuffle_target_builders=False,
165+
target_area=(1.5, 1.5),
166+
physics_timestep=_PHYSICS_TIMESTEP,
167+
control_timestep=_CONTROL_TIMESTEP,
168+
)
169+
170+
return composer.Environment(time_limit=30,
171+
task=task,
172+
random_state=random_state,
173+
strip_singleton_obs_buffer_dim=True)

‎dm_control/locomotion/props/__init__.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,3 +15,4 @@
1515
"""Props for Locomotion tasks."""
1616

1717
from dm_control.locomotion.props.target_sphere import TargetSphere
18+
from dm_control.locomotion.props.target_sphere import TargetSphereTwoTouch

0 commit comments

Comments
 (0)