[Question] How can I log Q-values of DQN using custom callback from stable baselines 3 in Tensorboard?

[gglsmm]

❓ Question

Hello,

I am trying to log Q-values using custom callback, but I am new in this field and not sure the code below is the correct way to do it.

class CustomLoggingCallback(BaseCallback):
    def __init__(self, verbose=1):
        super(CustomLoggingCallback, self).__init__(verbose)

    def _on_step(self) -> bool:
        # Log Q-values
        obs = th.tensor(self.locals["replay_buffer"].observations[-1], device=self.model.device).float()
        q_values = self.model.q_net(obs)
        avg_q_values = q_values.mean().item()
        self.logger.record('q_values/mean', avg_q_values)

        return True

Checklist

• [X] I have checked that there is no similar issue in the repo
• [X] I have read the documentation
• [X] If code there is, it is minimal and working
• [X] If code there is, it is formatted using the markdown code blocks for both code and stack traces.

[araffin]

Probably a duplicate of https://github.com/DLR-RM/stable-baselines3/issues/568 and https://github.com/DLR-RM/stable-baselines3/issues/308

[gglsmm]

Thank you for your response. in #568, I assume it is for after training.

My question is I want to log Q-values during the training process. For example, we observe loss or episode reward mean during the training on Tensorboard and I want to see Q-values if possible. See attached it is for Pong game. I am just not sure if this code is right.

[araffin]

I assume it is for after training.

why?

See attached it is for Pong game. I am just not sure if this code is right.

it looks fine, although you should use the code example I linked and you might want to call logger.dump() if the values are not logged often enough for you.

[gglsmm]

I assume it is for after training.

why?

I collect Q-values during training, so my callback is in model.learn(), but in this code #568, I assumed it collects Q-values after loading a trained model, if not it is my misunderstanding.

Thank you for your feedback.