def observation_space(self):
    """Returns the env's observation space in a Gym interface."""
    if self._dm_suite:
      return gym.spaces.Box(
          shape=self._env.observation_spec().shape,
          dtype=self._env.observation_spec().dtype,
          low=float('-inf'),
          high=float('+inf'),
      )
    else:
      return self._env.observation_space
Returns the env's observation space in a Gym interface.
trax/rl/task.py