Is there a way to log the weight changes of neural networks?

RLlib
1

Hello,

I would like to visualize how the weights are changing as a neural network continues its training. I’m prototyping this with RLlib DQN in simple Gym environments. Weights and Biases provides watch() method that logs the weight and gradients, but I’m having difficulty finding out how to integrate that into Ray Tune and RLlib.

Thanks in advance.

2

Hey, you could add the weights inside the on_train_result callback (override the rllib/agents/callbacks.py::DefaultCallbacks class and implement on_train_result method). In there is a result arg, which is a dict to which you can add arbitrary data (e.g. your weights, which you can get in that method via trainer.get_policy().get_weights()).

3

When I tried this I was able only to get the value with shape = 1 using W&B:

image
image760×330 9.61 KB

There is the way I tried to logged it:

class MyCallbacks(DefaultCallbacks):
    def on_train_result(self, trainer, result: dict, **kwargs):
        for k, v in trainer.get_policy().get_weights().items():
            result["FCC/{}".format(k)] = v

How can I log also other weights other than value_out/bias?
It seems that there is skipped all other arrays.
I also tried with v.flatten() version.

4

A bit of a tangent but worth a shot:

how did you make MyCallbacks compatible with W&B? Were you using tune or RLlib Trainer directly?

I am trying to make the logs from callbacks inheriting from DefaultCallbacks appear in W&B but can’t get that to work. I am using tune.run to train (as opposed to RLlib Trainer)

5

Sven I think this callback is so useful that deserves an example in the docs!