🔒 [TorchRL] What's wrong with the way I am saving my model?

[feral nova]

https://paste.pythondiscord.com/J5UQ

A question about TorchRL:

checkpoint = torch.load('ppo_model.pth')
actor_net.load_state_dict(checkpoint['actor_net_state_dict'])
value_net.load_state_dict(checkpoint['value_net_state_dict'])
ProbabilisticActor_policy_module.load_state_dict(checkpoint['probabilistic_actor_state_dict'])
scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
Adam_optimizer_used.load_state_dict(checkpoint['optimizer_state_dict'])
GAE_advantage_module.load_state_dict(checkpoint['gae_state_dict'])
maximum_average_reward = checkpoint['maximum_reward_tensor']

For some reason the average reward decreases after I load the model. I saved the states of Actor Network, Value Network, ProbabilisticActor state, Cosine Annealing Learning Rate Scheduler state, Adam optimiser state, Generalised Advantage Estimator State and maximum average reward number.

Here is where the saving happens:

torch.save({
  'actor_net_state_dict': actor_net.state_dict(),
  'value_net_state_dict': value_net.state_dict(),
  'probabilistic_actor_state_dict': ProbabilisticActor_policy_module.state_dict(),
  'optimizer_state_dict': Adam_optimizer_used.state_dict(),
  'scheduler_state_dict': scheduler.state_dict(),
  'gae_state_dict': GAE_advantage_module.state_dict(),
  'maximum_reward_tensor': maximum_average_reward_tensor
}, 'ppo_model.pth')

Do I need to save

1. sub-batch ReplayBuffer state
2. SyncDataCollector experience collector class state
3. ClipPPOLoss class state
4. Do I need to save the model before .step() method of my Cosine Annealing LR Scheduler or after?

Here are the average reward graphs for 1st and 2nd learning session.

[lavish verge]

@feral nova

Python help channel opened

Remember to:

• Ask your Python question, not if you can ask or if there's an expert who can help.
• Show a code sample as text (rather than a screenshot) and the error message, if you've got one.
• Explain what you expect to happen and what actually happens.

:warning: Do not pip install anything that isn't related to your question, especially if asked to over DMs.

Closes after a period of inactivity, or when you send !close.

[feral nova]

Am I not saving some important state which causes the agent performance to degrade inbetween sessions?

[lavish verge]

@feral nova

Python help channel closed

This help channel has been closed and it's no longer possible to send messages here. If your question wasn't answered, feel free to create a new post in #1035199133436354600. To maximize your chances of getting a response, check out this guide on asking good questions.