updated custom tasks to new api

This commit is contained in:
Fabian
2023-01-12 17:22:45 +01:00
parent fbe3ef4a4b
commit ed724046f3
23 changed files with 249 additions and 363 deletions
+86 -58
View File
@@ -1,18 +1,17 @@
import logging
import re
import uuid
from collections.abc import MutableMapping
from copy import deepcopy
from math import ceil
from typing import Iterable, Type, Union
import gym
import gymnasium as gym
import numpy as np
from gym.envs.registration import register, registry
from gym.utils import seeding
from gymnasium.envs.registration import register, registry
try:
from dm_control import suite, manipulation
from shimmy.dm_control_compatibility import EnvType
except ImportError:
pass
@@ -83,15 +82,20 @@ def make(env_id: str, seed: int, **kwargs):
if framework == 'metaworld':
# MetaWorld environment
env = make_metaworld(env_id, seed, **kwargs)
elif framework == 'dmc':
# DeepMind Control environment
env = make_dmc(env_id, seed, **kwargs)
# elif framework == 'dmc':
# Deprecated: With shimmy gym now has native support for deepmind envs
# # DeepMind Control environment
# env = make_dmc(env_id, seed, **kwargs)
else:
env = make_gym(env_id, seed, **kwargs)
np_random, _ = seeding.np_random(seed)
env.np_random = np_random
# env.seed(seed)
# try:
env.reset(seed=seed)
# except TypeError:
# # Support for older gym envs that do not have seeding
# # env.seed(seed)
# np_random, _ = seeding.np_random(seed)
# env.np_random = np_random
env.action_space.seed(seed)
env.observation_space.seed(seed)
@@ -161,7 +165,7 @@ def make_bb(
traj_gen_kwargs['action_dim'] = traj_gen_kwargs.get('action_dim', np.prod(env.action_space.shape).item())
if black_box_kwargs.get('duration') is None:
black_box_kwargs['duration'] = env.spec.max_episode_steps * env.dt
black_box_kwargs['duration'] = get_env_duration(env)
if phase_kwargs.get('tau') is None:
phase_kwargs['tau'] = black_box_kwargs['duration']
@@ -180,6 +184,24 @@ def make_bb(
return bb_env
def get_env_duration(env: gym.Env):
try:
# TODO Remove if this is in the compatibility class
duration = env.spec.max_episode_steps * env.dt
except (AttributeError, TypeError) as e:
logging.error(f'Attributes env.spec.max_episode_steps and env.dt are not available. '
f'Assuming you are using dm_control. Please make sure you have ran '
f'"pip install shimmy[dm_control]" for that.')
if env.env_type is EnvType.COMPOSER:
max_episode_steps = ceil(env.unwrapped._time_limit / env.dt)
elif env.env_type is EnvType.RL_CONTROL:
max_episode_steps = int(env.unwrapped._step_limit)
else:
raise e
duration = max_episode_steps * env.control_timestep()
return duration
def make_bb_env_helper(**kwargs):
"""
Helper function for registering a black box gym environment.
@@ -229,52 +251,53 @@ def make_bb_env_helper(**kwargs):
basis_kwargs=basis_kwargs, **kwargs, seed=seed)
def make_dmc(
env_id: str,
seed: int = None,
visualize_reward: bool = True,
time_limit: Union[None, float] = None,
**kwargs
):
if not re.match(r"\w+-\w+", env_id):
raise ValueError("env_id does not have the following structure: 'domain_name-task_name'")
domain_name, task_name = env_id.split("-")
if task_name.endswith("_vision"):
# TODO
raise ValueError("The vision interface for manipulation tasks is currently not supported.")
if (domain_name, task_name) not in suite.ALL_TASKS and task_name not in manipulation.ALL:
raise ValueError(f'Specified domain "{domain_name}" and task "{task_name}" combination does not exist.')
# env_id = f'dmc_{domain_name}_{task_name}_{seed}-v1'
gym_id = uuid.uuid4().hex + '-v1'
task_kwargs = {'random': seed}
if time_limit is not None:
task_kwargs['time_limit'] = time_limit
# create task
# Accessing private attribute because DMC does not expose time_limit or step_limit.
# Only the current time_step/time as well as the control_timestep can be accessed.
if domain_name == "manipulation":
env = manipulation.load(environment_name=task_name, seed=seed)
max_episode_steps = ceil(env._time_limit / env.control_timestep())
else:
env = suite.load(domain_name=domain_name, task_name=task_name, task_kwargs=task_kwargs,
visualize_reward=visualize_reward, environment_kwargs=kwargs)
max_episode_steps = int(env._step_limit)
register(
id=gym_id,
entry_point='fancy_gym.dmc.dmc_wrapper:DMCWrapper',
kwargs={'env': lambda: env},
max_episode_steps=max_episode_steps,
)
env = gym.make(gym_id)
env.seed(seed)
return env
# Deprecated: With shimmy gym now has native support for deepmind envs
# def make_dmc(
# env_id: str,
# seed: int = None,
# visualize_reward: bool = True,
# time_limit: Union[None, float] = None,
# **kwargs
# ):
# if not re.match(r"\w+-\w+", env_id):
# raise ValueError("env_id does not have the following structure: 'domain_name-task_name'")
# domain_name, task_name = env_id.split("-")
#
# if task_name.endswith("_vision"):
# # TODO
# raise ValueError("The vision interface for manipulation tasks is currently not supported.")
#
# if (domain_name, task_name) not in suite.ALL_TASKS and task_name not in manipulation.ALL:
# raise ValueError(f'Specified domain "{domain_name}" and task "{task_name}" combination does not exist.')
#
# # env_id = f'dmc_{domain_name}_{task_name}_{seed}-v1'
# gym_id = uuid.uuid4().hex + '-v1'
#
# task_kwargs = {'random': seed}
# if time_limit is not None:
# task_kwargs['time_limit'] = time_limit
#
# # create task
# # Accessing private attribute because DMC does not expose time_limit or step_limit.
# # Only the current time_step/time as well as the control_timestep can be accessed.
# if domain_name == "manipulation":
# env = manipulation.load(environment_name=task_name, seed=seed)
# max_episode_steps = ceil(env._time_limit / env.control_timestep())
# else:
# env = suite.load(domain_name=domain_name, task_name=task_name, task_kwargs=task_kwargs,
# visualize_reward=visualize_reward, environment_kwargs=kwargs)
# max_episode_steps = int(env._step_limit)
#
# register(
# id=gym_id,
# entry_point='fancy_gym.dmc.dmc_wrapper:DMCWrapper',
# kwargs={'env': lambda: env},
# max_episode_steps=max_episode_steps,
# )
#
# env = gym.make(gym_id)
# env.seed(seed)
# return env
def make_metaworld(env_id: str, seed: int, **kwargs):
@@ -288,12 +311,17 @@ def make_metaworld(env_id: str, seed: int, **kwargs):
# New argument to use global seeding
_env.seeded_rand_vec = True
max_episode_steps = _env.max_path_length
# TODO remove this as soon as there is support for the new API
_env = gym.wrappers.EnvCompatibility(_env)
gym_id = uuid.uuid4().hex + '-v1'
register(
id=gym_id,
entry_point=lambda: _env,
max_episode_steps=_env.max_path_length,
max_episode_steps=max_episode_steps,
)
# TODO enable checker when the incorrect dtype of obs and observation space are fixed by metaworld
+2 -60
View File
@@ -1,45 +1,11 @@
"""
Adapted from: https://github.com/openai/gym/blob/907b1b20dd9ac0cba5803225059b9c6673702467/gym/wrappers/time_aware_observation.py
License: MIT
Copyright (c) 2016 OpenAI (https://openai.com)
Wrapper for adding time aware observations to environment observation.
"""
import gym
import gymnasium as gym
import numpy as np
from gym.spaces import Box
class TimeAwareObservation(gym.ObservationWrapper):
"""Augment the observation with the current time step in the episode.
The observation space of the wrapped environment is assumed to be a flat :class:`Box`.
In particular, pixel observations are not supported. This wrapper will append the current timestep
within the current episode to the observation.
Example:
>>> import gym
>>> env = gym.make('CartPole-v1')
>>> env = TimeAwareObservation(env)
>>> env.reset()
array([ 0.03810719, 0.03522411, 0.02231044, -0.01088205, 0. ])
>>> env.step(env.action_space.sample())[0]
array([ 0.03881167, -0.16021058, 0.0220928 , 0.28875574, 1. ])
"""
class TimeAwareObservation(gym.wrappers.TimeAwareObservation):
def __init__(self, env: gym.Env):
"""Initialize :class:`TimeAwareObservation` that requires an environment with a flat :class:`Box`
observation space.
Args:
env: The environment to apply the wrapper
"""
super().__init__(env)
assert isinstance(env.observation_space, Box)
low = np.append(self.observation_space.low, 0.0)
high = np.append(self.observation_space.high, 1.0)
self.observation_space = Box(low, high, dtype=self.observation_space.dtype)
self.t = 0
self._max_episode_steps = env.spec.max_episode_steps
def observation(self, observation):
@@ -52,27 +18,3 @@ class TimeAwareObservation(gym.ObservationWrapper):
The observation with the time step appended to
"""
return np.append(observation, self.t / self._max_episode_steps)
def step(self, action):
"""Steps through the environment, incrementing the time step.
Args:
action: The action to take
Returns:
The environment's step using the action.
"""
self.t += 1
return super().step(action)
def reset(self, **kwargs):
"""Reset the environment setting the time to zero.
Args:
**kwargs: Kwargs to apply to env.reset()
Returns:
The reset environment
"""
self.t = 0
return super().reset(**kwargs)