1234567891011121314151617181920212223242526272829303132333435 |
- import logging
- import uuid
- from ray.rllib.examples.env.random_env import RandomEnv
- from ray.rllib.utils.annotations import override
- logger = logging.getLogger(__name__)
- class MemoryLeakingEnv(RandomEnv):
- """An env that leaks very little memory.
- Useful for proving that our memory-leak tests can catch the
- slightest leaks.
- """
- def __init__(self, config=None):
- super().__init__(config)
- self._leak = {}
- self._steps_after_reset = 0
- @override(RandomEnv)
- def reset(self, *, seed=None, options=None):
- self._steps_after_reset = 0
- return super().reset(seed=seed, options=options)
- @override(RandomEnv)
- def step(self, action):
- self._steps_after_reset += 1
- # Only leak once an episode.
- if self._steps_after_reset == 2:
- self._leak[uuid.uuid4().hex.upper()] = 1
- return super().step(action)
|