<gh_stars>1-10
# This file is a part of PDDLRL project.
# Copyright (c) 2020 <NAME> (<EMAIL>)
# Copyright (c) 2021 <NAME> (<EMAIL>, <EMAIL>), IBM Corporation

from acme.wrappers import base

import dm_env


class BoundedEpisodes(base.EnvironmentWrapper):

    def __init__(self, environment, max_epsiode_length):
        self.max_epsiode_length = max_epsiode_length
        self._count = 0
        super().__init__(environment)

    def reset(self):
        self._count = 0
        return self._environment.reset()

    def step(self, action):
        self._count += 1
        timestep = self._environment.step(action)
        if self._count >= self.max_epsiode_length:
            timestep = timestep._replace(step_type=dm_env.StepType.LAST)

        return timestep
