PPO rollouts not terminating with `done == True`

I am using a custom environment, and I’ve already checked the following:

from stable_baselines3.common.env_checker import check_env

env = CustomEnv(arg1, ...)
# It will check your custom environment and output additional warnings if needed
check_env(env)

But the PPO algo is continuing to step(action) when done == True (when the state is no longer in the bounds).

This is how I am interfacing with the algorithm:

class Agent:

    def __init__(self, environment, name, net_arch=[100, 100], n_env=1, n_steps=10000):

        # vectorise environment
        self.environment = environment
        check_env(self.environment)
        venv = DummyVecEnv([lambda: environment]*n_env)

        # load or create model
        assert(isinstance(name, str))
        self.name = name
        try:
            self.model = PPO.load(self.name, venv)
        except:
            self.model = PPO(
                'MlpPolicy', 
                venv,
                use_sde=True,
                sde_sample_freq=5,
                gae_lambda=0.9,
                learning_rate=1e-2,
                verbose=1,
                policy_kwargs=dict(net_arch=net_arch),
                n_steps=n_steps
            )

    def train(self, time_steps):

        # learn and save
        self.model.learn(total_timesteps=time_steps)
        self.model.save(self.name)

    def evaluate(self):

        # simulate
        obs = self.environment.reset()
        while True:
            action, _ = self.model.predict(obs, deterministic=True)
            obs, rew, done, _ = self.environment.step(action)
            if done: break

        # plot
        self.environment.system.plot(fname='{}.pdf'.format(self.name))

[cisprague]

@Miffyli, thanks. But, after more thought, I actually need the agent to stay within the state-space bounds during training, because some my underlying code requires that. How can I enforce this?

[Miffyli]

So, in this case, the reward itself should be enough to eventually enforce that I want the agent to stay within the bounds, right?

If I understood this right (you desire agent to avoid specific situations), then yes, a correct reward and/or terminal states should be able to teach agent to avoid those situations.