Gym Wrapper for Habitat environment compatible with tf-agents

habitat_wrapper.py

class HabitatWrapper(gym_wrapper.GymWrapper):
    # GymWrapper cannot handle all types of gym Spaces, and the action
    # space in Habitat (at least for the PointNav task) is a gym.Dict
    # (supported) with str keys and habitat.EmptySpace values
    # (unsupported). Since the action space is really a discrete space,
    # we'll update gym_env.action_space temporarily to be Discrete
    # so that GymWrapper can create an action spec.

    def __init__(self,
                 gym_env,
                 discount=1.0,
                 spec_dtype_map=None,
                 match_obs_space_dtype=True,
                 auto_reset=True,
                 simplify_box_bounds=True):
        original_action_space = gym_env.action_space
        gym_env.action_space = original_action_space.actions_select
        super(HabitatWrapper, self).__init__(
            gym_env, discount, spec_dtype_map, match_obs_space_dtype,
            auto_reset, simplify_box_bounds)
        gym_env.action_space = original_action_space