Updated train_on_episode_end

[LucaVendruscolo]

Hello, I had a few issues with git so I deleted my old fork #319 and started again. I made all the requested changes and did some general testing on all the algorithms to confirm it worked with them too. Let me know if there's anything else you'd like me to add/change.

Thanks so much for the help!

[belerico]

Hi @LucaVendruscolo, do you have some empirical evidence that everything is working? Some plots for example

[LucaVendruscolo]

I will run the tests and take some screenshots now

[LucaVendruscolo]

The log files: logs.zip

grey line is train_on_episode_end: false and blue line is train_on_episode_end: true
running:
python sheeprl.py exp=dreamer_v3 env=gym env.id=CartPole-v1 env.num_envs=4