Get the current step number in a gym.Env

Is there a way to access the current step number of a gym.Env from inside its step method?

I'm using a model from stable_baselines3 and want to terminate the env when N steps have been taken.

You can make your own TerminatingEnv class that wraps an enviroment and takes into account counting and terminating if the count is too high internally, so you can use it like any other environment while still having the terminatiion feature. env.step_counter will tell you how many steps have been done since the reset of the enviroment:

import gym

class TerminatingEnv(gym.Env):
    def __init__(self, env, max_steps):
        self.env = env
        self.max_steps = max_steps
        self.step_counter = 0

    def step(self, action):
        
        obs, reward, terminated, truncated, info = self.env.step(action)
        self.step_counter += 1
        if self.step_counter >= self.max_steps:
            return obs, reward, True, True, info
        else:
            return obs, reward, terminated, truncated, info
    
    def reset(self):
        self.step_counter = 0
        return self.env.reset()
    
    def render(self):
        return self.env.render()
    
    def close(self):
        return self.env.close()
    
    def seed(self, seed=None):
        return self.env.seed(seed)
    

env0 = gym.make('CartPole-v0')
env = TerminatingEnv(env0, 5)
while True:
    action = 1
    env.reset()
    for i in range(100):
        obs, reward, terminated, truncated, info = env.step(action)
        print(env.step_counter)
        if terminated:
            break

The output is:

1
2
3
4
5
1
2
3
4
5
1

And so on...

In older versions, self.env.step returns 4 values and not 5, so maybe you will have to adapt that.