Files
Gymnasium/gym/wrappers/time_aware_observation.py
Andrea PIERRÉ e913bc81b8 Improve pre-commit workflow (#2602)
* feat: add `isort` to `pre-commit`

* ci: skip `__init__.py` file for `isort`

* ci: make `isort` mandatory in lint pipeline

* docs: add a section on Git hooks

* ci: check isort diff

* fix: isort from master branch

* docs: add pre-commit badge

* ci: update black + bandit versions

* feat: add PR template

* refactor: PR template

* ci: remove bandit

* docs: add Black badge

* ci: try to remove all `|| true` statements

* ci: remove lint_python job

- Remove `lint_python` CI job
- Move `pyupgrade` job to `pre-commit` workflow

* fix: avoid messing with typing

* docs: add a note on running `pre-cpmmit` manually

* ci: apply `pre-commit` to the whole codebase
2022-03-31 15:50:38 -04:00

36 lines
1005 B
Python

from typing import Optional
import numpy as np
from gym import ObservationWrapper
from gym.spaces import Box
class TimeAwareObservation(ObservationWrapper):
r"""Augment the observation with current time step in the trajectory.
.. note::
Currently it only works with one-dimensional observation space. It doesn't
support pixel observation space yet.
"""
def __init__(self, env):
super().__init__(env)
assert isinstance(env.observation_space, Box)
assert env.observation_space.dtype == np.float32
low = np.append(self.observation_space.low, 0.0)
high = np.append(self.observation_space.high, np.inf)
self.observation_space = Box(low, high, dtype=np.float32)
def observation(self, observation):
return np.append(observation, self.t)
def step(self, action):
self.t += 1
return super().step(action)
def reset(self, **kwargs):
self.t = 0
return super().reset(**kwargs)