dmp env wrappers initial
This commit is contained in:
Maximilian Huettenrauch
@@ -0,0 +1,171 @@
|
||||
import gym
|
||||
from gym.error import (AlreadyPendingCallError, NoAsyncCallError)
|
||||
from gym.vector.utils import concatenate, create_empty_array
|
||||
from gym.vector.async_vector_env import AsyncState
|
||||
import numpy as np
|
||||
import multiprocessing as mp
|
||||
from copy import deepcopy
|
||||
import sys
|
||||
|
||||
|
||||
class DmpAsyncVectorEnv(gym.vector.AsyncVectorEnv):
|
||||
def __init__(self, env_fns, n_samples, observation_space=None, action_space=None,
|
||||
shared_memory=True, copy=True, context=None, daemon=True, worker=None):
|
||||
super(DmpAsyncVectorEnv, self).__init__(env_fns,
|
||||
observation_space=observation_space,
|
||||
action_space=action_space,
|
||||
shared_memory=shared_memory,
|
||||
copy=copy,
|
||||
context=context,
|
||||
daemon=daemon,
|
||||
worker=worker)
|
||||
|
||||
# we need to overwrite the number of samples as we may sample more than num_envs
|
||||
self.observations = create_empty_array(self.single_observation_space,
|
||||
n=n_samples,
|
||||
fn=np.zeros)
|
||||
|
||||
def __call__(self, params):
|
||||
return self.rollout(params)
|
||||
|
||||
def rollout_async(self, actions):
|
||||
"""
|
||||
Parameters
|
||||
----------
|
||||
actions : iterable of samples from `action_space`
|
||||
List of actions.
|
||||
"""
|
||||
self._assert_is_running()
|
||||
if self._state != AsyncState.DEFAULT:
|
||||
raise AlreadyPendingCallError('Calling `rollout_async` while waiting '
|
||||
'for a pending call to `{0}` to complete.'.format(
|
||||
self._state.value), self._state.value)
|
||||
|
||||
# split_actions = np.array_split(actions, self.num_envs)
|
||||
actions = np.atleast_2d(actions)
|
||||
split_actions = np.array_split(actions, np.minimum(len(actions), self.num_envs))
|
||||
for pipe, action in zip(self.parent_pipes, split_actions):
|
||||
pipe.send(('rollout', action))
|
||||
for pipe in self.parent_pipes[len(split_actions):]:
|
||||
pipe.send(('idle', None))
|
||||
self._state = AsyncState.WAITING_ROLLOUT
|
||||
|
||||
def rollout_wait(self, timeout=None):
|
||||
"""
|
||||
Parameters
|
||||
----------
|
||||
timeout : int or float, optional
|
||||
Number of seconds before the call to `step_wait` times out. If
|
||||
`None`, the call to `step_wait` never times out.
|
||||
|
||||
Returns
|
||||
-------
|
||||
observations : sample from `observation_space`
|
||||
A batch of observations from the vectorized environment.
|
||||
|
||||
rewards : `np.ndarray` instance (dtype `np.float_`)
|
||||
A vector of rewards from the vectorized environment.
|
||||
|
||||
dones : `np.ndarray` instance (dtype `np.bool_`)
|
||||
A vector whose entries indicate whether the episode has ended.
|
||||
|
||||
infos : list of dict
|
||||
A list of auxiliary diagnostic information.
|
||||
"""
|
||||
self._assert_is_running()
|
||||
if self._state != AsyncState.WAITING_ROLLOUT:
|
||||
raise NoAsyncCallError('Calling `rollout_wait` without any prior call '
|
||||
'to `rollout_async`.', AsyncState.WAITING_ROLLOUT.value)
|
||||
|
||||
if not self._poll(timeout):
|
||||
self._state = AsyncState.DEFAULT
|
||||
raise mp.TimeoutError('The call to `rollout_wait` has timed out after '
|
||||
'{0} second{1}.'.format(timeout, 's' if timeout > 1 else ''))
|
||||
|
||||
results, successes = zip(*[pipe.recv() for pipe in self.parent_pipes])
|
||||
results = [r for r in results if r is not None]
|
||||
self._raise_if_errors(successes)
|
||||
self._state = AsyncState.DEFAULT
|
||||
|
||||
observations_list, rewards, dones, infos = [_flatten_list(r) for r in zip(*results)]
|
||||
|
||||
# if not self.shared_memory:
|
||||
# self.observations = concatenate(observations_list, self.observations,
|
||||
# self.single_observation_space)
|
||||
|
||||
# return (deepcopy(self.observations) if self.copy else self.observations,
|
||||
# np.array(rewards), np.array(dones, dtype=np.bool_), infos)
|
||||
|
||||
return np.array(rewards)
|
||||
|
||||
def rollout(self, actions):
|
||||
self.rollout_async(actions)
|
||||
return self.rollout_wait()
|
||||
|
||||
|
||||
def _worker(index, env_fn, pipe, parent_pipe, shared_memory, error_queue):
|
||||
assert shared_memory is None
|
||||
env = env_fn()
|
||||
parent_pipe.close()
|
||||
try:
|
||||
while True:
|
||||
command, data = pipe.recv()
|
||||
if command == 'reset':
|
||||
observation = env.reset()
|
||||
pipe.send((observation, True))
|
||||
elif command == 'step':
|
||||
observation, reward, done, info = env.step(data)
|
||||
if done:
|
||||
observation = env.reset()
|
||||
pipe.send(((observation, reward, done, info), True))
|
||||
elif command == 'rollout':
|
||||
observations = []
|
||||
rewards = []
|
||||
dones = []
|
||||
infos = []
|
||||
for d in data:
|
||||
env.reset()
|
||||
observation, reward, done, info = env.step(d)
|
||||
observations.append(observation)
|
||||
rewards.append(reward)
|
||||
dones.append(done)
|
||||
infos.append(info)
|
||||
pipe.send(((observations, rewards, dones, infos), (True, ) * len(rewards)))
|
||||
elif command == 'seed':
|
||||
env.seed(data)
|
||||
pipe.send((None, True))
|
||||
elif command == 'close':
|
||||
pipe.send((None, True))
|
||||
break
|
||||
elif command == 'idle':
|
||||
pipe.send((None, True))
|
||||
elif command == '_check_observation_space':
|
||||
pipe.send((data == env.observation_space, True))
|
||||
else:
|
||||
raise RuntimeError('Received unknown command `{0}`. Must '
|
||||
'be one of {`reset`, `step`, `seed`, `close`, '
|
||||
'`_check_observation_space`}.'.format(command))
|
||||
except (KeyboardInterrupt, Exception):
|
||||
error_queue.put((index,) + sys.exc_info()[:2])
|
||||
pipe.send((None, False))
|
||||
finally:
|
||||
env.close()
|
||||
|
||||
|
||||
def _flatten_obs(obs):
|
||||
assert isinstance(obs, (list, tuple))
|
||||
assert len(obs) > 0
|
||||
|
||||
if isinstance(obs[0], dict):
|
||||
keys = obs[0].keys()
|
||||
return {k: np.stack([o[k] for o in obs]) for k in keys}
|
||||
else:
|
||||
return np.stack(obs)
|
||||
|
||||
|
||||
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_]
|
||||
@@ -0,0 +1,102 @@
|
||||
from mp_lib.phase import ExpDecayPhaseGenerator
|
||||
from mp_lib.basis import DMPBasisGenerator
|
||||
from mp_lib import dmps
|
||||
import numpy as np
|
||||
import gym
|
||||
|
||||
|
||||
class DmpEnvWrapperBase(gym.Wrapper):
|
||||
def __init__(self, env, num_dof, num_basis, duration=1, dt=0.01, learn_goal=False):
|
||||
super(DmpEnvWrapperBase, self).__init__(env)
|
||||
self.num_dof = num_dof
|
||||
self.num_basis = num_basis
|
||||
self.dim = num_dof * num_basis
|
||||
if learn_goal:
|
||||
self.dim += num_dof
|
||||
self.learn_goal = True
|
||||
duration = duration # seconds
|
||||
time_steps = int(duration / dt)
|
||||
self.t = np.linspace(0, duration, time_steps)
|
||||
|
||||
phase_generator = ExpDecayPhaseGenerator(alpha_phase=5, duration=duration)
|
||||
basis_generator = DMPBasisGenerator(phase_generator, duration=duration, num_basis=self.num_basis)
|
||||
|
||||
self.dmp = dmps.DMP(num_dof=num_dof,
|
||||
basis_generator=basis_generator,
|
||||
phase_generator=phase_generator,
|
||||
num_time_steps=time_steps,
|
||||
dt=dt
|
||||
)
|
||||
|
||||
self.dmp.dmp_start_pos = env.start_pos.reshape((1, num_dof))
|
||||
|
||||
dmp_weights = np.zeros((num_basis, num_dof))
|
||||
dmp_goal_pos = np.zeros(num_dof)
|
||||
|
||||
self.dmp.set_weights(dmp_weights, dmp_goal_pos)
|
||||
|
||||
def goal_and_weights(self, params):
|
||||
if len(params.shape) > 1:
|
||||
assert params.shape[1] == self.dim
|
||||
else:
|
||||
assert len(params) == self.dim
|
||||
params = np.reshape(params, [1, self.dim])
|
||||
|
||||
if self.learn_goal:
|
||||
goal_pos = params[0, -self.num_dof:]
|
||||
weight_matrix = np.reshape(params[:, :-self.num_dof], [self.num_basis, self.num_dof])
|
||||
else:
|
||||
goal_pos = None
|
||||
weight_matrix = np.reshape(params, [self.num_basis, self.num_dof])
|
||||
|
||||
return goal_pos, weight_matrix
|
||||
|
||||
def step(self, action, render=False):
|
||||
""" overwrite step function where action now is the weights and possible goal position"""
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class DmpEnvWrapperAngle(DmpEnvWrapperBase):
|
||||
def step(self, action, render=False):
|
||||
goal_pos, weight_matrix = self.goal_and_weights(action)
|
||||
self.dmp.set_weights(weight_matrix, goal_pos)
|
||||
trajectory, velocities = self.dmp.reference_trajectory(self.t)
|
||||
|
||||
rews = []
|
||||
|
||||
for t, traj in enumerate(trajectory):
|
||||
obs, rew, done, info = self.env.step(traj)
|
||||
rews.append(rew)
|
||||
if render:
|
||||
self.env.render(mode="human")
|
||||
if done:
|
||||
break
|
||||
|
||||
reward = np.sum(rews)
|
||||
done = True
|
||||
info = {}
|
||||
|
||||
return obs, reward, done, info
|
||||
|
||||
|
||||
class DmpEnvWrapperVel(DmpEnvWrapperBase):
|
||||
def step(self, action, render=False):
|
||||
goal_pos, weight_matrix = self.goal_and_weights(action)
|
||||
weight_matrix *= 50
|
||||
self.dmp.set_weights(weight_matrix, goal_pos)
|
||||
trajectory, velocities = self.dmp.reference_trajectory(self.t)
|
||||
|
||||
rews = []
|
||||
|
||||
for t, vel in enumerate(velocities):
|
||||
obs, rew, done, info = self.env.step(vel)
|
||||
rews.append(rew)
|
||||
if render:
|
||||
self.env.render(mode="human")
|
||||
if done:
|
||||
break
|
||||
|
||||
reward = np.sum(rews)
|
||||
info = {}
|
||||
|
||||
return obs, reward, done, info
|
||||
Reference in New Issue
Block a user