if task_env.ModalityTypes.IMAGE in self._modality_types:
      obs_shapes[task_env.ModalityTypes.IMAGE] = gym.spaces.Box(
          low=0, high=255, shape=(self._output_size, self._output_size, 3))