import gym from mujoco_py import MujocoException from gym.spaces import Dict, Box import numpy as np from copy import deepcopy import logging def update_obs_space(env, delta): spaces = env.observation_space.spaces.copy() for key, shape in delta.items(): spaces[key] = Box(-np.inf, np.inf, shape, np.float32) return Dict(spaces) class NumpyArrayRewardWrapper(gym.RewardWrapper): """ Convenience wrapper that casts rewards to the multiagent format (numpy array of shape (n_agents,)) """ def __init__(self, env): super().__init__(env) def reward(self, rew): return np.zeros((self.unwrapped.n_agents,)) + rew class DiscretizeActionWrapper(gym.ActionWrapper): ''' Take a Box action and convert it to a MultiDiscrete Action through quantization Args: action_key: (string) action to discretize nbuckets: (int) number of discrete actions per dimension. It should be odd such that actions centered around 0 will have the middle action be 0. ''' def __init__(self, env, action_key, nbuckets=11): super().__init__(env) self.action_key = action_key self.discrete_to_continuous_act_map = [] for i, ac_space in enumerate(self.action_space.spaces[action_key].spaces): assert isinstance(ac_space, Box) action_map = np.array([np.linspace(low, high, nbuckets) for low, high in zip(ac_space.low, ac_space.high)]) _nbuckets = np.ones((len(action_map))) * nbuckets self.action_space.spaces[action_key].spaces[i] = gym.spaces.MultiDiscrete(_nbuckets) self.discrete_to_continuous_act_map.append(action_map) self.discrete_to_continuous_act_map = np.array(self.discrete_to_continuous_act_map) def action(self, action): action = deepcopy(action) ac = action[self.action_key] # helper variables for indexing the discrete-to-continuous action map agent_idxs = np.tile(np.arange(ac.shape[0])[:, None], ac.shape[1]) ac_idxs = np.tile(np.arange(ac.shape[1]), ac.shape[0]).reshape(ac.shape) action[self.action_key] = self.discrete_to_continuous_act_map[agent_idxs, ac_idxs, ac] return action class DiscardMujocoExceptionEpisodes(gym.Wrapper): ''' Catches Mujoco Exceptions. Sends signal to discard Episode. ''' def __init__(self, env): super().__init__(env) self.episode_error = False def step(self, action): assert not self.episode_error, "Won't Continue Episode After Mujoco Exception -- \ Please discard episode and reset. If info['discard_episode'] is True the episode\ should be discarded" try: obs, rew, done, info = self.env.step(action) info['discard_episode'] = False except MujocoException as e: self.episode_error = True # Done is set to False such that rollout workers do not accidently send data in # the event that timelimit is up in the same step as an error occured. obs, rew, done, info = {}, 0.0, False, {'discard_episode': True} logging.info(str(e)) logging.info("Encountered Mujoco Exception During Environment Step.\ Reset Episode Required") return obs, rew, done, info def reset(self): try: obs = self.env.reset() except MujocoException: logging.info("Encountered Mujoco Exception During Environment Reset.\ Trying Reset Again") obs = self.reset() self.episode_error = False return obs class MaskActionWrapper(gym.Wrapper): ''' For a boolean action, sets it to zero given a mask from the previous step. For example you could mask the grab action based on whether you can see the box Args: action_key (string): key in action dictionary to be masked mask_keys (string): keys in observation dictionary with which to mask. The shape of the concatenation of the masks (along the 1st dimension) should exactly match that of action_key ''' def __init__(self, env, action_key, mask_keys): super().__init__(env) self.action_key = action_key self.mask_keys = mask_keys def reset(self): self.prev_obs = self.env.reset() return deepcopy(self.prev_obs) def step(self, action): mask = np.concatenate([self.prev_obs[k] for k in self.mask_keys], -1) action[self.action_key] = np.logical_and(action[self.action_key], mask) self.prev_obs, rew, done, info = self.env.step(action) return deepcopy(self.prev_obs), rew, done, info class AddConstantObservationsWrapper(gym.ObservationWrapper): ''' Adds new constant observations to the environment. Args: new_obs: Dictionary with the new observations. ''' def __init__(self, env, new_obs): super().__init__(env) self.new_obs = new_obs for obs_key in self.new_obs: assert obs_key not in self.observation_space.spaces, ( f'Observation key {obs_key} exists in original observation space') if type(self.new_obs[obs_key]) in [list, tuple]: self.new_obs[obs_key] = np.array(self.new_obs[obs_key]) shape = self.new_obs[obs_key].shape self.observation_space = update_obs_space(self, {obs_key: shape}) def observation(self, obs): for key, val in self.new_obs.items(): obs[key] = val return obs class SpoofEntityWrapper(gym.ObservationWrapper): ''' Add extra entities along entity dimension such that shapes can match between environments with differing number of entities. This is meant to be used after SplitObservations and SelectKeysWrapper. This will also add masks that are 1 except along the new columns (which could be used by fully observed value function) Args: total_n_entities (int): total number of entities after spoofing (including spoofed ones) keys (list): observation keys with which to add entities along the second dimension mask_keys (list): mask keys with which to add columns. ''' def __init__(self, env, total_n_entities, keys, mask_keys): super().__init__(env) self.total_n_entities = total_n_entities self.keys = keys self.mask_keys = mask_keys for key in self.keys + self.mask_keys: shape = list(self.observation_space.spaces[key].shape) shape[1] = total_n_entities self.observation_space = update_obs_space(self, {key: shape}) for key in self.mask_keys: shape = list(self.observation_space.spaces[key].shape) self.observation_space = update_obs_space(self, {key + '_spoof': shape}) def observation(self, obs): for key in self.keys: n_to_spoof = self.total_n_entities - obs[key].shape[1] if n_to_spoof > 0: obs[key] = np.concatenate([obs[key], np.zeros((obs[key].shape[0], n_to_spoof, obs[key].shape[-1]))], 1) for key in self.mask_keys: n_to_spoof = self.total_n_entities - obs[key].shape[1] obs[key + '_spoof'] = np.concatenate([np.ones_like(obs[key]), np.zeros((obs[key].shape[0], n_to_spoof))], -1) if n_to_spoof > 0: obs[key] = np.concatenate([obs[key], np.zeros((obs[key].shape[0], n_to_spoof))], -1) return obs class ConcatenateObsWrapper(gym.ObservationWrapper): ''' Group multiple observations under the same key in the observation dictionary. Args: obs_groups: dict of {key_to_save: [keys to concat]} ''' def __init__(self, env, obs_groups): super().__init__(env) self.obs_groups = obs_groups for key_to_save, keys_to_concat in obs_groups.items(): assert np.all([np.array(self.observation_space.spaces[keys_to_concat[0]].shape[:-1]) == np.array(self.observation_space.spaces[k].shape[:-1]) for k in keys_to_concat]), \ f"Spaces were {[(k, v) for k, v in self.observation_space.spaces.items() if k in keys_to_concat]}" new_last_dim = sum([self.observation_space.spaces[k].shape[-1] for k in keys_to_concat]) new_shape = list(self.observation_space.spaces[keys_to_concat[0]].shape[:-1]) + [new_last_dim] self.observation_space = update_obs_space(self, {key_to_save: new_shape}) def observation(self, obs): for key_to_save, keys_to_concat in self.obs_groups.items(): obs[key_to_save] = np.concatenate([obs[k] for k in keys_to_concat], -1) return obs