updates

alr_envs/utils/mp_env_async_sampler.py

@@ -0,0 +1,82 @@
+import gym
+from gym.vector.async_vector_env import AsyncVectorEnv
+import numpy as np
+from _collections import defaultdict
+
+
+def make_env(env_id, rank, seed=0):
+    env = gym.make(env_id)
+    env.seed(seed + rank)
+    return lambda: env
+
+
+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 % repeat != 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 AlrMpEnvSampler:
+    """
+    An asynchronous sampler for 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):
+        self.num_envs = num_envs
+        self.env = AsyncVectorEnv([make_env(env_id, seed, i) 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:
+            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]
+
+
+if __name__ == "__main__":
+    env_name = "alr_envs:HoleReacherDMP-v0"
+    n_cpu = 8
+    dim = 30
+    n_samples = 20
+
+    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)

alr_envs/utils/wrapper/dmp_wrapper.py

@@ -11,8 +11,9 @@ class DmpWrapper(MPWrapper):
 
     def __init__(self, env: gym.Env, num_dof: int, num_basis: int, start_pos: np.ndarray = None,
                  final_pos: np.ndarray = None, duration: int = 1, alpha_phase: float = 2., dt: float = None,
-                 learn_goal: bool = False, post_traj_time: float = 0., policy_type: str = None,
-                 weights_scale: float = 1., goal_scale: float = 1., bandwidth_factor: float = 3.):
+                 learn_goal: bool = False, return_to_start: bool = False, post_traj_time: float = 0.,
+                 weights_scale: float = 1., goal_scale: float = 1., bandwidth_factor: float = 3.,
+                 policy_type: str = None):
 
         """
         This Wrapper generates a trajectory based on a DMP and will only return episodic performances.
@@ -34,8 +35,13 @@ class DmpWrapper(MPWrapper):
         self.learn_goal = learn_goal
         dt = env.dt if hasattr(env, "dt") else dt
         assert dt is not None
-        start_pos = env.start_pos if hasattr(env, "start_pos") else start_pos
+        start_pos = start_pos if start_pos is not None else env.start_pos if hasattr(env, "start_pos") else None
         assert start_pos is not None
+        if learn_goal:
+            final_pos = np.zeros_like(start_pos)  # arbitrary, will be learned
+        else:
+            final_pos = final_pos if final_pos is not None else start_pos if return_to_start else None
+        assert final_pos is not None
         self.t = np.linspace(0, duration, int(duration / dt))
         self.goal_scale = goal_scale
 

alr_envs/utils/wrapper/mp_wrapper.py

@@ -46,8 +46,9 @@ class MPWrapper(gym.Wrapper, ABC):
         rewards = []
         dones = []
         infos = []
-        for p, c in zip(params, contexts):
-            self.configure(c)
+        # for p, c in zip(params, contexts):
+        for p in params:
+            # self.configure(c)
             ob, reward, done, info = self.step(p)
             obs.append(ob)
             rewards.append(reward)