fix time aware observation wrapper issue and add try catch for rendering dmc

alr_envs/utils/make_env_helpers.py

@@ -28,7 +28,7 @@ from alr_envs.black_box.factory.controller_factory import get_controller
 from alr_envs.black_box.factory.phase_generator_factory import get_phase_generator
 from alr_envs.black_box.factory.trajectory_generator_factory import get_trajectory_generator
 from alr_envs.black_box.raw_interface_wrapper import RawInterfaceWrapper
-from alr_envs.black_box.time_aware_observation import TimeAwareObservation
+from alr_envs.utils.time_aware_observation import TimeAwareObservation
 from alr_envs.utils.utils import nested_update
 
 
@@ -148,11 +148,10 @@ def make_bb(
     if learn_sub_trajs and do_replanning:
         raise ValueError('Cannot used sub-trajectory learning and replanning together.')
 
-    if learn_sub_trajs or do_replanning:
-        # add time_step observation when replanning
-        if not any(issubclass(w, TimeAwareObservation) for w in kwargs['wrappers']):
-            # Add as first wrapper in order to alter observation
-            kwargs['wrappers'].insert(0, TimeAwareObservation)
+    # add time_step observation when replanning
+    if (learn_sub_trajs or do_replanning) and not any(issubclass(w, TimeAwareObservation) for w in kwargs['wrappers']):
+        # Add as first wrapper in order to alter observation
+        kwargs['wrappers'].insert(0, TimeAwareObservation)
 
     env = _make_wrapped_env(env_id=env_id, wrappers=wrappers, seed=seed, **kwargs)
 

alr_envs/utils/time_aware_observation.py

@@ -0,0 +1,78 @@
+"""
+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 numpy as np
+
+import gym
+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.        ])
+    """
+
+    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, np.inf)
+        self.observation_space = Box(low, high, dtype=self.observation_space.dtype)
+        self.t = 0
+
+    def observation(self, observation):
+        """Adds to the observation with the current time step.
+
+        Args:
+            observation: The observation to add the time step to
+
+        Returns:
+            The observation with the time step appended to
+        """
+        return np.append(observation, self.t)
+
+    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)