fixed OpenAI fetch tasks; added nicer imports

alr_envs/utils/__init__.py

@@ -37,7 +37,6 @@ def make(
         episode_length = 250 if domain_name == "manipulation" else 1000
 
     max_episode_steps = (episode_length + frame_skip - 1) // frame_skip
-
     if env_id not in gym.envs.registry.env_specs:
         task_kwargs = {'random': seed}
         # if seed is not None:
@@ -46,7 +45,7 @@ def make(
             task_kwargs['time_limit'] = time_limit
         register(
             id=env_id,
-            entry_point='alr_envs.utils.dmc2gym_wrapper:DMCWrapper',
+            entry_point='alr_envs.utils.dmc_wrapper:DMCWrapper',
             kwargs=dict(
                 domain_name=domain_name,
                 task_name=task_name,

alr_envs/utils/dmc_wrapper.py

@@ -33,11 +33,14 @@ def _spec_to_box(spec):
 
 
 def _flatten_obs(obs: collections.MutableMapping):
-    # obs_pieces = []
-    # for v in obs.values():
-    #     flat = np.array([v]) if np.isscalar(v) else v.ravel()
-    #     obs_pieces.append(flat)
-    # return np.concatenate(obs_pieces, axis=0)
+    """
+    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.')
@@ -52,19 +55,19 @@ def _flatten_obs(obs: collections.MutableMapping):
 class DMCWrapper(core.Env):
     def __init__(
             self,
-            domain_name,
-            task_name,
-            task_kwargs={},
-            visualize_reward=True,
-            from_pixels=False,
-            height=84,
-            width=84,
-            camera_id=0,
-            frame_skip=1,
-            environment_kwargs=None,
-            channels_first=True
+            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 behaviour'
+        assert 'random' in task_kwargs, 'Please specify a seed for deterministic behavior.'
         self._from_pixels = from_pixels
         self._height = height
         self._width = width
@@ -74,7 +77,7 @@ class DMCWrapper(core.Env):
 
         # create task
         if domain_name == "manipulation":
-            assert not from_pixels, \
+            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:
@@ -169,11 +172,12 @@ class DMCWrapper(core.Env):
         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
-            camera_id = camera_id or self._camera_id
             return self._env.physics.render(height=height, width=width, camera_id=camera_id)
 
         elif mode == 'human':
@@ -184,7 +188,8 @@ class DMCWrapper(core.Env):
                 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)
+                                           width=self._env.physics.model.vis.global_.offwidth,
+                                           camera_id=camera_id)
             self.viewer.imshow(img)
             return self.viewer.isopen
 

alr_envs/utils/make_env_helpers.py

@@ -2,13 +2,14 @@ import logging
 from typing import Iterable, List, Type, Union
 
 import gym
+import numpy as np
 
-from mp_env_api.interface_wrappers.mp_env_wrapper import MPEnvWrapper
+from mp_env_api import MPEnvWrapper
 from mp_env_api.mp_wrappers.detpmp_wrapper import DetPMPWrapper
 from mp_env_api.mp_wrappers.dmp_wrapper import DmpWrapper
 
 
-def make_env_rank(env_id: str, seed: int, rank: int = 0, **kwargs):
+def make_env_rank(env_id: str, seed: int, rank: int = 0, return_callable=True, **kwargs):
     """
     TODO: Do we need this?
     Generate a callable to create a new gym environment with a given seed.
@@ -22,11 +23,16 @@ def make_env_rank(env_id: str, seed: int, rank: int = 0, **kwargs):
         env_id: name of the environment
         seed: seed for deterministic behaviour
         rank: environment rank for deterministic over multiple seeds behaviour
+        return_callable: If True returns a callable to create the environment instead of the environment itself.
 
     Returns:
 
     """
-    return lambda: make_env(env_id, seed + rank, **kwargs)
+
+    def f():
+        return make_env(env_id, seed + rank, **kwargs)
+
+    return f if return_callable else f()
 
 
 def make_env(env_id: str, seed, **kwargs):
@@ -103,6 +109,9 @@ def make_dmp_env(env_id: str, wrappers: Iterable, seed=1, mp_kwargs={}, **kwargs
     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"))
+
     return DmpWrapper(_env, **mp_kwargs)
 
 
@@ -120,6 +129,9 @@ def make_detpmp_env(env_id: str, wrappers: Iterable, seed=1, mp_kwargs={}, **kwa
     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"))
+
     return DetPMPWrapper(_env, **mp_kwargs)
 
 
@@ -185,5 +197,12 @@ def verify_time_limit(mp_time_limit: Union[None, float], env_time_limit: Union[N
     """
     if mp_time_limit is not None and env_time_limit is not None:
         assert mp_time_limit == env_time_limit, \
-            f"The manually specified 'time_limit' of {env_time_limit}s does not match " \
+            f"The specified 'time_limit' of {env_time_limit}s does not match " \
             f"the duration of {mp_time_limit}s for the MP."
+
+
+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 " \
+        f"the action space of {action_shape} the base environments"

alr_envs/utils/utils.py

@@ -15,8 +15,7 @@ def angle_normalize(x, type="deg"):
     if type not in ["deg", "rad"]: raise ValueError(f"Invalid type {type}. Choose one of 'deg' or 'rad'.")
 
     if type == "deg":
-        x = np.deg2rad(x) # x * pi / 180
+        x = np.deg2rad(x)  # x * pi / 180
 
     two_pi = 2 * np.pi
     return x - two_pi * np.floor((x + np.pi) / two_pi)
-