Episode mean reward is not properly logged on tensorboard when using SAC

I’m training an agent on a custom environment using SAC. The environment is wrapped in a Monitor, which is wrapped in a DummyVecEnv, which is wrapped in a VecNormalize, with norm_reward = True.

This is the tensorboard graph for the episode mean reward:

No smoothing	0.9 Smoothing

As you can see, the graph has some weird loops. For example, at around 170k steps or 450k steps.

---

Edit: Training is conducted in epochs of 50k steps.

Program starts by calling ./start.sh.

start.sh

#!/bin/bash

while [ "$?" -eq 0 ]; do
	python3 main.py
done

main.py

import os.path

from my_custom_env import MyCustomEnv

from stable_baselines3 import SAC
from stable_baselines3.common.callbacks import BaseCallback
from stable_baselines3.common.monitor import Monitor
from stable_baselines3.common.vec_env import DummyVecEnv, VecNormalize

class SaveCheckpoint(BaseCallback):
	def __init__(self, save_freq, verbose = 0):
		super(SaveCheckpoint, self).__init__(verbose)
		self.save_freq = save_freq

	def _on_step(self):
		if self.num_timesteps % self.save_freq == 0:
			self.model.save("model.zip")
			self.training_env.save("stats.pkl")

		return True


if __name__ == '__main__':

	# inits
	env = DummyVecEnv([lambda: Monitor(MyCustomEnv())])
	model = None

	# load recent checkpoint
	if os.path.isfile("model.zip") and os.path.isfile("stats.pkl"):
		env = VecNormalize.load("stats.pkl", env)
		env.reset()
		model = SAC.load("model.zip", env)
	else:
		env = VecNormalize(env)
		model = SAC('MlpPolicy', env, verbose = 1, tensorboard_log = ".")

	# replay buffer
	if os.path.isfile("replay_buffer.pkl"):
		model.load_replay_buffer("replay_buffer.pkl")

	# train
	model.learn(50000,
		callback = SaveCheckpoint(10000),
		log_interval = 1,
		reset_num_timesteps = False
	)

	# save replay buffer
	model.save_replay_buffer(".")

	env.close()

> pip3 freeze | grep 'stable-baselines3'
stable-baselines3==0.7.0a1

[siferati]

Updated the post with more detailed info

[siferati]

Yup, that seems to have been the problem. Thanks!