from gym.wrappers import TimeLimit


class MultiTimeLimit(TimeLimit):
    def step(self, action):
        observation, reward, done, info = self.env.step(action)
        self._elapsed_steps += 1
        if self._elapsed_steps >= self._max_episode_steps:
            for agent_idx in range(self.env.config.num_agents):
                info[agent_idx]["TimeLimit.truncated"] = not done[agent_idx]
            done = [True] * self.num_agents
        return observation, reward, done, info
