torch.save(
{
"model": agent.state_dict(),
"optimizer": optimizer.state_dict(),
"scheduler": scheduler.state_dict(),
"episode": settings.EPISODE_LEN,
"synthetic_episode": settings.SYNTHETIC_EPISODE_LEN,
"epoch": epoch
},
f=settings.SAVE_PATH)
def test(epoch):
pass
if __name__ == "__main__":
for epoch in range(start_epoch, settings.EPOCH):
train(epoch)
test(epoch)
scheduler.step()
settings.step_episode(
inc=1
) # WARNING: If you are using multi-batch, the episode len while training should be 1. which means we do not rendering. If you want to do render, make sure the model in one batch is the same, or just keep the batch size === 1.
settings.step_synthetic_episode(inc=1)
# Save
settings.RESUME = True
settings.dump_yaml("settings/ycb.yaml")
