| 84 | |
| 85 | |
| 86 | class LimitLength(gym.Wrapper): |
| 87 | def __init__(self, env, k): |
| 88 | gym.Wrapper.__init__(self, env) |
| 89 | self.k = k |
| 90 | |
| 91 | def reset(self): |
| 92 | # This assumes that reset() will really reset the env. |
| 93 | # If the underlying env tries to be smart about reset |
| 94 | # (e.g. end-of-life), the assumption doesn't hold. |
| 95 | ob = self.env.reset() |
| 96 | self.cnt = 0 |
| 97 | return ob |
| 98 | |
| 99 | def step(self, action): |
| 100 | ob, r, done, info = self.env.step(action) |
| 101 | self.cnt += 1 |
| 102 | if self.cnt == self.k: |
| 103 | done = True |
| 104 | return ob, r, done, info |