improved project structure and exposed methods
This commit is contained in:
ottofabian
@@ -47,7 +47,7 @@ def make_dmc(
|
||||
task_kwargs['time_limit'] = time_limit
|
||||
register(
|
||||
id=env_id,
|
||||
entry_point='alr_envs.utils.dmc_wrapper:DMCWrapper',
|
||||
entry_point='alr_envs.dmc.dmc_wrapper:DMCWrapper',
|
||||
kwargs=dict(
|
||||
domain_name=domain_name,
|
||||
task_name=task_name,
|
||||
|
||||
@@ -1,206 +0,0 @@
|
||||
# Adopted from: https://github.com/denisyarats/dmc2gym/blob/master/dmc2gym/wrappers.py
|
||||
# License: MIT
|
||||
# Copyright (c) 2020 Denis Yarats
|
||||
import collections
|
||||
from typing import Any, Dict, Tuple
|
||||
|
||||
import numpy as np
|
||||
from dm_control import manipulation, suite
|
||||
from dm_env import specs
|
||||
from gym import core, spaces
|
||||
|
||||
|
||||
def _spec_to_box(spec):
|
||||
def extract_min_max(s):
|
||||
assert s.dtype == np.float64 or s.dtype == np.float32, f"Only float64 and float32 types are allowed, instead {s.dtype} was found"
|
||||
dim = int(np.prod(s.shape))
|
||||
if type(s) == specs.Array:
|
||||
bound = np.inf * np.ones(dim, dtype=np.float32)
|
||||
return -bound, bound
|
||||
elif type(s) == specs.BoundedArray:
|
||||
zeros = np.zeros(dim, dtype=np.float32)
|
||||
return s.minimum + zeros, s.maximum + zeros
|
||||
|
||||
mins, maxs = [], []
|
||||
for s in spec:
|
||||
mn, mx = extract_min_max(s)
|
||||
mins.append(mn)
|
||||
maxs.append(mx)
|
||||
low = np.concatenate(mins, axis=0)
|
||||
high = np.concatenate(maxs, axis=0)
|
||||
assert low.shape == high.shape
|
||||
return spaces.Box(low, high, dtype=np.float32)
|
||||
|
||||
|
||||
def _flatten_obs(obs: collections.MutableMapping):
|
||||
"""
|
||||
Flattens an observation of type MutableMapping, e.g. a dict to a 1D array.
|
||||
Args:
|
||||
obs: observation to flatten
|
||||
|
||||
Returns: 1D array of observation
|
||||
|
||||
"""
|
||||
|
||||
if not isinstance(obs, collections.MutableMapping):
|
||||
raise ValueError(f'Requires dict-like observations structure. {type(obs)} found.')
|
||||
|
||||
# Keep key order consistent for non OrderedDicts
|
||||
keys = obs.keys() if isinstance(obs, collections.OrderedDict) else sorted(obs.keys())
|
||||
|
||||
obs_vals = [np.array([obs[key]]) if np.isscalar(obs[key]) else obs[key].ravel() for key in keys]
|
||||
return np.concatenate(obs_vals)
|
||||
|
||||
|
||||
class DMCWrapper(core.Env):
|
||||
def __init__(
|
||||
self,
|
||||
domain_name: str,
|
||||
task_name: str,
|
||||
task_kwargs: dict = {},
|
||||
visualize_reward: bool = True,
|
||||
from_pixels: bool = False,
|
||||
height: int = 84,
|
||||
width: int = 84,
|
||||
camera_id: int = 0,
|
||||
frame_skip: int = 1,
|
||||
environment_kwargs: dict = None,
|
||||
channels_first: bool = True
|
||||
):
|
||||
assert 'random' in task_kwargs, 'Please specify a seed for deterministic behavior.'
|
||||
self._from_pixels = from_pixels
|
||||
self._height = height
|
||||
self._width = width
|
||||
self._camera_id = camera_id
|
||||
self._frame_skip = frame_skip
|
||||
self._channels_first = channels_first
|
||||
|
||||
# create task
|
||||
if domain_name == "manipulation":
|
||||
assert not from_pixels and not task_name.endswith("_vision"), \
|
||||
"TODO: Vision interface for manipulation is different to suite and needs to be implemented"
|
||||
self._env = manipulation.load(environment_name=task_name, seed=task_kwargs['random'])
|
||||
else:
|
||||
self._env = suite.load(domain_name=domain_name, task_name=task_name, task_kwargs=task_kwargs,
|
||||
visualize_reward=visualize_reward, environment_kwargs=environment_kwargs)
|
||||
|
||||
# action and observation space
|
||||
self._action_space = _spec_to_box([self._env.action_spec()])
|
||||
self._observation_space = _spec_to_box(self._env.observation_spec().values())
|
||||
|
||||
self._last_state = None
|
||||
self.viewer = None
|
||||
|
||||
# set seed
|
||||
self.seed(seed=task_kwargs.get('random', 1))
|
||||
|
||||
def __getattr__(self, item):
|
||||
"""Propagate only non-existent properties to wrapped env."""
|
||||
if item.startswith('_'):
|
||||
raise AttributeError("attempted to get missing private attribute '{}'".format(item))
|
||||
if item in self.__dict__:
|
||||
return getattr(self, item)
|
||||
return getattr(self._env, item)
|
||||
|
||||
def _get_obs(self, time_step):
|
||||
if self._from_pixels:
|
||||
obs = self.render(
|
||||
mode="rgb_array",
|
||||
height=self._height,
|
||||
width=self._width,
|
||||
camera_id=self._camera_id
|
||||
)
|
||||
if self._channels_first:
|
||||
obs = obs.transpose(2, 0, 1).copy()
|
||||
else:
|
||||
obs = _flatten_obs(time_step.observation)
|
||||
return obs
|
||||
|
||||
@property
|
||||
def observation_space(self):
|
||||
return self._observation_space
|
||||
|
||||
@property
|
||||
def action_space(self):
|
||||
return self._action_space
|
||||
|
||||
@property
|
||||
def dt(self):
|
||||
return self._env.control_timestep() * self._frame_skip
|
||||
|
||||
@property
|
||||
def base_step_limit(self):
|
||||
"""
|
||||
Returns: max_episode_steps of the underlying DMC env
|
||||
|
||||
"""
|
||||
# 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.
|
||||
try:
|
||||
return (self._env._step_limit + self._frame_skip - 1) // self._frame_skip
|
||||
except AttributeError as e:
|
||||
return self._env._time_limit / self.dt
|
||||
|
||||
def seed(self, seed=None):
|
||||
self._action_space.seed(seed)
|
||||
self._observation_space.seed(seed)
|
||||
|
||||
def step(self, action) -> Tuple[np.ndarray, float, bool, Dict[str, Any]]:
|
||||
assert self._action_space.contains(action)
|
||||
reward = 0
|
||||
extra = {'internal_state': self._env.physics.get_state().copy()}
|
||||
|
||||
for _ in range(self._frame_skip):
|
||||
time_step = self._env.step(action)
|
||||
reward += time_step.reward or 0.
|
||||
done = time_step.last()
|
||||
if done:
|
||||
break
|
||||
|
||||
self._last_state = _flatten_obs(time_step.observation)
|
||||
obs = self._get_obs(time_step)
|
||||
extra['discount'] = time_step.discount
|
||||
return obs, reward, done, extra
|
||||
|
||||
def reset(self) -> np.ndarray:
|
||||
time_step = self._env.reset()
|
||||
self._last_state = _flatten_obs(time_step.observation)
|
||||
obs = self._get_obs(time_step)
|
||||
return obs
|
||||
|
||||
def render(self, mode='rgb_array', height=None, width=None, camera_id=0):
|
||||
if self._last_state is None:
|
||||
raise ValueError('Environment not ready to render. Call reset() first.')
|
||||
|
||||
camera_id = camera_id or self._camera_id
|
||||
|
||||
# assert mode == 'rgb_array', 'only support rgb_array mode, given %s' % mode
|
||||
if mode == "rgb_array":
|
||||
height = height or self._height
|
||||
width = width or self._width
|
||||
return self._env.physics.render(height=height, width=width, camera_id=camera_id)
|
||||
|
||||
elif mode == 'human':
|
||||
if self.viewer is None:
|
||||
# pylint: disable=import-outside-toplevel
|
||||
# pylint: disable=g-import-not-at-top
|
||||
from gym.envs.classic_control import rendering
|
||||
self.viewer = rendering.SimpleImageViewer()
|
||||
# Render max available buffer size. Larger is only possible by altering the XML.
|
||||
img = self._env.physics.render(height=self._env.physics.model.vis.global_.offheight,
|
||||
width=self._env.physics.model.vis.global_.offwidth,
|
||||
camera_id=camera_id)
|
||||
self.viewer.imshow(img)
|
||||
return self.viewer.isopen
|
||||
|
||||
def close(self):
|
||||
super().close()
|
||||
if self.viewer is not None and self.viewer.isopen:
|
||||
self.viewer.close()
|
||||
|
||||
@property
|
||||
def reward_range(self) -> Tuple[float, float]:
|
||||
reward_spec = self._env.reward_spec()
|
||||
if isinstance(reward_spec, specs.BoundedArray):
|
||||
return reward_spec.minimum, reward_spec.maximum
|
||||
return -float('inf'), float('inf')
|
||||
@@ -76,7 +76,7 @@ def make(env_id: str, seed, **kwargs):
|
||||
|
||||
else:
|
||||
# DMC
|
||||
from alr_envs.utils import make_dmc
|
||||
from alr_envs import make_dmc
|
||||
env = make_dmc(env_id, seed=seed, **kwargs)
|
||||
|
||||
assert env.base_step_limit == env.spec.max_episode_steps, \
|
||||
@@ -123,11 +123,11 @@ def make_dmp_env(env_id: str, wrappers: Iterable, seed=1, mp_kwargs={}, **kwargs
|
||||
Returns: DMP wrapped gym env
|
||||
|
||||
"""
|
||||
verify_time_limit(mp_kwargs.get("duration", None), kwargs.get("time_limit", None))
|
||||
_verify_time_limit(mp_kwargs.get("duration", None), kwargs.get("time_limit", None))
|
||||
|
||||
_env = _make_wrapped_env(env_id=env_id, wrappers=wrappers, seed=seed, **kwargs)
|
||||
|
||||
verify_dof(_env, mp_kwargs.get("num_dof"))
|
||||
_verify_dof(_env, mp_kwargs.get("num_dof"))
|
||||
|
||||
return DmpWrapper(_env, **mp_kwargs)
|
||||
|
||||
@@ -143,11 +143,11 @@ def make_detpmp_env(env_id: str, wrappers: Iterable, seed=1, mp_kwargs={}, **kwa
|
||||
Returns: DMP wrapped gym env
|
||||
|
||||
"""
|
||||
verify_time_limit(mp_kwargs.get("duration", None), kwargs.get("time_limit", None))
|
||||
_verify_time_limit(mp_kwargs.get("duration", None), kwargs.get("time_limit", None))
|
||||
|
||||
_env = _make_wrapped_env(env_id=env_id, wrappers=wrappers, seed=seed, **kwargs)
|
||||
|
||||
verify_dof(_env, mp_kwargs.get("num_dof"))
|
||||
_verify_dof(_env, mp_kwargs.get("num_dof"))
|
||||
|
||||
return DetPMPWrapper(_env, **mp_kwargs)
|
||||
|
||||
@@ -191,14 +191,7 @@ def make_detpmp_env_helper(**kwargs):
|
||||
mp_kwargs=kwargs.pop("mp_kwargs"), **kwargs)
|
||||
|
||||
|
||||
def make_contextual_env(env_id, context, seed, rank):
|
||||
env = make(env_id, seed + rank, context=context)
|
||||
# env = gym.make(env_id, context=context)
|
||||
# env.seed(seed + rank)
|
||||
return lambda: env
|
||||
|
||||
|
||||
def verify_time_limit(mp_time_limit: Union[None, float], env_time_limit: Union[None, float]):
|
||||
def _verify_time_limit(mp_time_limit: Union[None, float], env_time_limit: Union[None, float]):
|
||||
"""
|
||||
When using DMC check if a manually specified time limit matches the trajectory duration the MP receives.
|
||||
Mostly, the time_limit for DMC is not specified and the default values from DMC are taken.
|
||||
@@ -218,7 +211,7 @@ def verify_time_limit(mp_time_limit: Union[None, float], env_time_limit: Union[N
|
||||
f"the duration of {mp_time_limit}s for the MP."
|
||||
|
||||
|
||||
def verify_dof(base_env: gym.Env, dof: int):
|
||||
def _verify_dof(base_env: gym.Env, dof: int):
|
||||
action_shape = np.prod(base_env.action_space.shape)
|
||||
assert dof == action_shape, \
|
||||
f"The specified degrees of freedom ('num_dof') {dof} do not match " \
|
||||
|
||||
@@ -1,122 +0,0 @@
|
||||
import gym
|
||||
from gym.vector.async_vector_env import AsyncVectorEnv
|
||||
import numpy as np
|
||||
from _collections import defaultdict
|
||||
|
||||
from alr_envs.utils.make_env_helpers import make_rank
|
||||
|
||||
|
||||
def split_array(ary, size):
|
||||
n_samples = len(ary)
|
||||
if n_samples < size:
|
||||
tmp = np.zeros((size, ary.shape[1]))
|
||||
tmp[0:n_samples] = ary
|
||||
return [tmp]
|
||||
elif n_samples == size:
|
||||
return [ary]
|
||||
else:
|
||||
repeat = int(np.ceil(n_samples / size))
|
||||
split = [k * size for k in range(1, repeat)]
|
||||
sub_arys = np.split(ary, split)
|
||||
|
||||
if n_samples % size != 0:
|
||||
tmp = np.zeros_like(sub_arys[0])
|
||||
last = sub_arys[-1]
|
||||
tmp[0: len(last)] = last
|
||||
sub_arys[-1] = tmp
|
||||
|
||||
return sub_arys
|
||||
|
||||
|
||||
def _flatten_list(l):
|
||||
assert isinstance(l, (list, tuple))
|
||||
assert len(l) > 0
|
||||
assert all([len(l_) > 0 for l_ in l])
|
||||
|
||||
return [l__ for l_ in l for l__ in l_]
|
||||
|
||||
|
||||
class DummyDist:
|
||||
def __init__(self, dim):
|
||||
self.dim = dim
|
||||
|
||||
def sample(self, contexts):
|
||||
contexts = np.atleast_2d(contexts)
|
||||
n_samples = contexts.shape[0]
|
||||
return np.random.normal(size=(n_samples, self.dim)), contexts
|
||||
|
||||
|
||||
class AlrMpEnvSampler:
|
||||
"""
|
||||
An asynchronous sampler for non contextual MPWrapper environments. A sampler object can be called with a set of
|
||||
parameters and returns the corresponding final obs, rewards, dones and info dicts.
|
||||
"""
|
||||
|
||||
def __init__(self, env_id, num_envs, seed=0, **env_kwargs):
|
||||
self.num_envs = num_envs
|
||||
self.env = AsyncVectorEnv([make_rank(env_id, seed, i, **env_kwargs) for i in range(num_envs)])
|
||||
|
||||
def __call__(self, params):
|
||||
params = np.atleast_2d(params)
|
||||
n_samples = params.shape[0]
|
||||
split_params = split_array(params, self.num_envs)
|
||||
|
||||
vals = defaultdict(list)
|
||||
for p in split_params:
|
||||
self.env.reset()
|
||||
obs, reward, done, info = self.env.step(p)
|
||||
vals['obs'].append(obs)
|
||||
vals['reward'].append(reward)
|
||||
vals['done'].append(done)
|
||||
vals['info'].append(info)
|
||||
|
||||
# do not return values above threshold
|
||||
return np.vstack(vals['obs'])[:n_samples], np.hstack(vals['reward'])[:n_samples], \
|
||||
_flatten_list(vals['done'])[:n_samples], _flatten_list(vals['info'])[:n_samples]
|
||||
|
||||
|
||||
class AlrContextualMpEnvSampler:
|
||||
"""
|
||||
An asynchronous sampler for contextual MPWrapper environments. A sampler object can be called with a set of
|
||||
parameters and returns the corresponding final obs, rewards, dones and info dicts.
|
||||
"""
|
||||
|
||||
def __init__(self, env_id, num_envs, seed=0, **env_kwargs):
|
||||
self.num_envs = num_envs
|
||||
self.env = AsyncVectorEnv([make_env(env_id, seed, i, **env_kwargs) for i in range(num_envs)])
|
||||
|
||||
def __call__(self, dist, n_samples):
|
||||
repeat = int(np.ceil(n_samples / self.env.num_envs))
|
||||
vals = defaultdict(list)
|
||||
|
||||
obs = self.env.reset()
|
||||
for i in range(repeat):
|
||||
vals['obs'].append(obs)
|
||||
new_samples, new_contexts = dist.sample(obs)
|
||||
vals['new_samples'].append(new_samples)
|
||||
|
||||
obs, reward, done, info = self.env.step(new_samples)
|
||||
|
||||
vals['reward'].append(reward)
|
||||
vals['done'].append(done)
|
||||
vals['info'].append(info)
|
||||
|
||||
# do not return values above threshold
|
||||
return np.vstack(vals['new_samples'])[:n_samples], \
|
||||
np.vstack(vals['obs'])[:n_samples], np.hstack(vals['reward'])[:n_samples], \
|
||||
_flatten_list(vals['done'])[:n_samples], _flatten_list(vals['info'])[:n_samples]
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
env_name = "alr_envs:HoleReacherDetPMP-v1"
|
||||
n_cpu = 8
|
||||
dim = 25
|
||||
n_samples = 10
|
||||
|
||||
sampler = AlrMpEnvSampler(env_name, num_envs=n_cpu)
|
||||
|
||||
thetas = np.random.randn(n_samples, dim) # usually form a search distribution
|
||||
|
||||
_, rewards, __, ___ = sampler(thetas)
|
||||
|
||||
print(rewards)
|
||||
Reference in New Issue
Block a user