Search code examples
pythonreinforcement-learningopenai-gymstable-baselines

Stablebaselines3 logging reward with custom gym


I have this custom callback to log the reward in my custom vectorized environment, but the reward appears in console as always [0] and is not logged in tensorboard at all

class TensorboardCallback(BaseCallback):
    """
    Custom callback for plotting additional values in tensorboard.
    """

    def __init__(self, verbose=0):
        super(TensorboardCallback, self).__init__(verbose)

    def _on_step(self) -> bool:                
        self.logger.record('reward', self.training_env.get_attr('total_reward'))
        return True

And this is part of the main function

model = PPO(
        "MlpPolicy", env,
        learning_rate=3e-4,
        policy_kwargs=policy_kwargs,
        verbose=1,

# as the environment is not serializable, we need to set a new instance of the environment
loaded_model = model = PPO.load("model", env=env)
loaded_model.set_env(env)

# and continue training
loaded_model.learn(1e+6, callback=TensorboardCallback())
        tensorboard_log="./tensorboard/")

Solution

  • You need to add [0] as indexing,

    so where you wrote self.logger.record('reward', self.training_env.get_attr('total_reward')) you just need to index with self.logger.record('reward', self.training_env.get_attr ('total_reward')[0])

    class TensorboardCallback(BaseCallback):
        """
        Custom callback for plotting additional values in tensorboard.
        """
    
        def __init__(self, verbose=0):
            super(TensorboardCallback, self).__init__(verbose)
    
        def _on_step(self) -> bool:                
            self.logger.record('reward', self.training_env.get_attr('total_reward')[0])
    
            return True