wrappers.py

import gym
from gym import Wrapper
from gym.spaces import flatten, flatten_space
from minedojo.sim.spaces import Discrete, MultiDiscrete, Box, Tuple, Dict, Text
import numpy as np
from copy import copy

class RemoveTextObservation(Wrapper):
    '''wrapper to remove text spaces from the observation space.'''

    def __init__(self, env):
        super().__init__(env)
        self._removed_paths = []

        def _remove_text_spaces(space, path: list=[]):
            '''recursive function to remove text spaces from the observation space.'''

            if isinstance(space, (Discrete, MultiDiscrete, Box)):
                return space
            elif isinstance(space, Tuple):
                tuple = []
                for i, s in enumerate(space.spaces):
                    sub_space = _remove_text_spaces(s, path + [i])
                    if sub_space != None:
                        tuple.append(sub_space)
                if len(tuple) == 0:
                    self._removed_paths.append(path)
                    return None
                else:
                    return Tuple(tuple)
            elif isinstance(space, Dict):
                dict = {}
                for k, v in space.spaces.items():
                    sub_space = _remove_text_spaces(v, path + [k])
                    if sub_space != None:
                        dict[k] = sub_space
                if len(dict) == 0:
                    self._removed_paths.append(path)
                    return None
                else:
                    return Dict(dict)
            elif isinstance(space, Text):
                self._removed_paths.append(path)
                return None
            else:
                raise ValueError(f"Unsupported space type: {type(space)}")
            
        self.observation_space = _remove_text_spaces(env.observation_space)
    
    def _remove_text_values(self, obs):
        '''function to remove text values according to the removed spaces from the observation.'''

        def _remove_text_value(obs, path: list[list]=[]):
            '''recursive function to remove text values according to the removed spaces from the observation.'''
            if path in self._removed_paths:
                return None
            else:
                if isinstance(obs, (list, tuple)):
                    return type(obs)([_remove_text_value(o, path + [i]) for i, o in enumerate(obs)])
                elif isinstance(obs, dict):
                    return {k: _remove_text_value(v, path + [k]) for k, v in obs.items()}
                else:
                    return obs
        
        obs = _remove_text_value(obs)
        return obs

    def reset(self, **kwargs):
        obs = self.env.reset(**kwargs)
        return self._remove_text_values(obs)

    def step(self, action):
        obs, reward, done, info = self.env.step(action)
        return self._remove_text_values(obs), reward, done, info


class SplittedFlattenObservation(Wrapper):
    '''wrapper that splits out certain high-dimensional observation spaces and have other spaces flattened.'''
    def __init__(self, env):
        super().__init__(env)
        # self.split_keys = ('rgb', 'voxels', 'lidar')
        self.split_keys = ('rgb',)
        self._flatten_spaces = copy(env.observation_space)
        non_flatten_spaces = {}
        for key in self.split_keys:
            if key not in self._flatten_spaces.keys():
                continue
            non_flatten_spaces[key] = self._flatten_spaces[key]
            del self._flatten_spaces[key]
        self._before_flatten_spaces = copy(self._flatten_spaces)
        self._flatten_spaces = flatten_space(self._flatten_spaces)
        non_flatten_spaces = Dict(non_flatten_spaces)
        self.observation_space = Dict({"flatten": self._flatten_spaces, "nonflatten": non_flatten_spaces})
        self.action_space = env.action_space
    
    def _split_flatten_obs(self, obs):
        '''function to split out certain high-dimensional observation spaces and have other spaces flattened.'''
        flatten_obs = copy(obs)
        non_flatten_obs = {}
        for key in self.split_keys:
            non_flatten_obs[key] = flatten_obs[key]
            del flatten_obs[key]
        flatten_obs = flatten(self._before_flatten_spaces, flatten_obs)
        return {"flatten": flatten_obs, "nonflatten": non_flatten_obs}
        
    def reset(self, **kwargs):
        obs = self.env.reset(**kwargs)
        obs = self._split_flatten_obs(obs)
        return obs

    def step(self, action):
        obs, reward, done, info = self.env.step(action)
        obs = self._split_flatten_obs(obs)
        return obs, reward, done, info