def main(unused_argv):
if not FLAGS.davis_path:
raise ValueError(
'You must download and extract the DAVIS dataset and pass a path '
'a path to the videos, e.g.: /tmp/DAVIS/JPEGImages/480p')
for i, difficulty in enumerate(['easy', 'medium', 'hard']):
for j, (domain, task) in enumerate([('ball_in_cup', 'catch'),
('cartpole', 'swingup'),
('cheetah', 'run'),
('finger', 'spin'),
('reacher', 'easy'),
('walker', 'walk')]):
env = suite.load(domain,
task,
difficulty,
background_dataset_path=FLAGS.davis_path)
# Get the first frame.
time_step = env.reset()
frame = time_step.observation['pixels'][:, :, 0:3]
# Save the first frame.
try:
os.mkdir(FLAGS.output_dir)
except OSError:
pass
filepath = os.path.join(FLAGS.output_dir, f'{i:02d}-{j:02d}.jpg')
image = PIL.Image.fromarray(frame)
image.save(filepath)
print(f'Saved results to {FLAGS.output_dir}')
def test_suite_load_with_difficulty(self, difficulty, mock_dm_suite,
mock_pixels):
domain_name = 'cartpole'
task_name = 'balance'
suite.load(domain_name,
task_name,
difficulty,
background_dataset_path=DAVIS_PATH)
mock_dm_suite.load.assert_called_with(domain_name,
task_name,
environment_kwargs=None,
task_kwargs=None,
visualize_reward=False)
mock_pixels.Wrapper.assert_called_with(mock.ANY,
observation_key='pixels',
pixels_only=True,
render_kwargs={'camera_id': 0})
def _load_dm_env(domain_name,
task_name,
pixels,
action_repeat,
max_episode_steps=None,
obs_type='pixels',
distractor=False):
"""Load a Deepmind control suite environment."""
try:
if not pixels:
env = suite_dm_control.load(domain_name=domain_name,
task_name=task_name)
if action_repeat > 1:
env = wrappers.ActionRepeat(env, action_repeat)
else:
def wrap_repeat(env):
return ActionRepeatDMWrapper(env, action_repeat)
camera_id = 2 if domain_name == 'quadruped' else 0
pixels_only = obs_type == 'pixels'
if distractor:
render_kwargs = dict(width=84, height=84, camera_id=camera_id)
env = distractor_suite.load(
domain_name,
task_name,
difficulty='hard',
dynamic=False,
background_dataset_path='DAVIS/JPEGImages/480p/',
task_kwargs={},
environment_kwargs={},
render_kwargs=render_kwargs,
visualize_reward=False,
env_state_wrappers=[wrap_repeat])
# env = wrap_repeat(env)
# env = suite.wrappers.pixels.Wrapper(
# env,
# pixels_only=pixels_only,
# render_kwargs=render_kwargs,
# observation_key=obs_type)
env = dm_control_wrapper.DmControlWrapper(env, render_kwargs)
else:
env = suite_dm_control.load_pixels(
domain_name=domain_name,
task_name=task_name,
render_kwargs=dict(width=84,
height=84,
camera_id=camera_id),
env_state_wrappers=[wrap_repeat],
observation_key=obs_type,
pixels_only=pixels_only)
if action_repeat > 1 and max_episode_steps is not None:
# Shorten episode length.
max_episode_steps = (max_episode_steps + action_repeat -
1) // action_repeat
env = wrappers.TimeLimit(env, max_episode_steps)
return env
except ValueError as e:
logging.warning(
'cannot instantiate dm env: domain_name=%s, task_name=%s',
domain_name, task_name)
logging.warning('Supported domains and tasks: %s',
str({
key: list(val.SUITE.keys())
for key, val in suite._DOMAINS.items()
})) # pylint: disable=protected-access
raise e
def load_pixels(
domain_name,
task_name,
observation_key='pixels',
pixels_only=True,
task_kwargs=None,
environment_kwargs=None,
visualize_reward=False,
render_kwargs=None,
env_wrappers=(),
camera_kwargs=None,
background_kwargs=None,
color_kwargs=None,
):
"""Returns an environment from a domain name, task name and optional settings.
Args:
domain_name: A string containing the name of a domain.
task_name: A string containing the name of a task.
observation_key: Optional custom string specifying the pixel observation's
key in the `OrderedDict` of observations. Defaults to 'pixels'.
pixels_only: If True (default), the original set of 'state' observations
returned by the wrapped environment will be discarded, and the
`OrderedDict` of observations will only contain pixels. If False, the
`OrderedDict` will contain the original observations as well as the pixel
observations.
task_kwargs: Optional `dict` of keyword arguments for the task.
environment_kwargs: Optional `dict` specifying keyword arguments for the
environment.
visualize_reward: Optional `bool`. If `True`, object colours in rendered
frames are set to indicate the reward at each step. Default `False`.
render_kwargs: Optional `dict` of keyword arguments for rendering.
env_wrappers: Iterable with references to wrapper classes to use on the
wrapped environment.
camera_kwargs: optional dict of camera distraction arguments
background_kwargs: optional dict of background distraction arguments
color_kwargs: optional dict of color distraction arguments
Returns:
The requested environment.
Raises:
ImportError: if dm_control module was not available.
"""
dm_env = suite.load(domain_name,
task_name,
task_kwargs=task_kwargs,
environment_kwargs=environment_kwargs,
visualize_reward=visualize_reward,
camera_kwargs=camera_kwargs,
background_kwargs=background_kwargs,
color_kwargs=color_kwargs,
pixels_only=pixels_only,
render_kwargs=render_kwargs,
pixels_observation_key=observation_key)
env = dm_control_wrapper.DmControlWrapper(dm_env, render_kwargs)
for wrapper in env_wrappers:
env = wrapper(env)
return env