Fixed save/load problem on dqn.py

[jmribeiro]

Saving and loading the DQN agent would not save/load four needed attributes:

• self.t
• self.optim_t
• self._cumulative_steps
• self.replay_buffer

This caused the agent to have different a performance when evaluated without killing the program vs when saving the agent, killing the program, resuming the program and loading the agent.

Fig 1 - Training without checkpoints (i.e. same program ran from start to finish)

Fig 2 - Training with checkpoint (i.e., program killed at every t steps and agents loaded from disk)

My proposed solution (working, but applied only to the DQN agent) was to add new save_snapshot and load_snapshot methods on the agent's class (without overwriting the original save and load methods, avoiding saving the replay buffer every time):

    def save_snapshot(self, dirname: str) -> None:
        self.save(dirname)
        torch.save(
            self.t, os.path.join(dirname, "t.pt")
        )
        torch.save(
            self.optim_t, os.path.join(dirname, "optim_t.pt")
        )
        torch.save(
            self._cumulative_steps, os.path.join(dirname, "_cumulative_steps.pt")
        )
        self.replay_buffer.save(
            os.path.join(dirname, "replay_buffer.pkl")
        )


    def load_snapshot(self, dirname: str) -> None:
        self.load(dirname)
        self.t = torch.load(
            os.path.join(dirname, "t.pt")
        )
        self.optim_t = torch.load(
            os.path.join(dirname, "optim_t.pt")
        )
        self._cumulative_steps = torch.load(
            os.path.join(dirname, "_cumulative_steps.pt")
        )
        self.replay_buffer.load(
            os.path.join(dirname, "replay_buffer.pkl")
        )

This change is working as intended, training is resumed properly after reloading the agent from disk:

Fig 3 - Training with checkpoint (New patch) (i.e., program killed at every t steps and agents loaded from disk)

[muupan]

/test

[pfn-ci-bot]

Successfully created a job for commit 8fc26f4:

• Dashboard for commit 8fc26f4

[jmribeiro]

/test

@muupan There seem to be a problem with tests on test_acer.py (unrelated to the changes)

[muupan]

Sorry for my delayed response! Some of CI including code lint is broken and should be fixed later. I see the changes by the PR is harmless besides potential lint warnings, which should also be fixed later if any. Thanks for your contribution!