Not clear how to debug

[cool-RR]

Description

I wanted to get a better idea of the values of variables on line 79 of iterated_matrix_game.py. I wanted to know what the observation is in that environment, since I couldn't understand it exactly from the way it was phrased in the paper. I stuck a jax.debug.breakpoint() in there, hoping that I could run pax.experiment and get the jdb prompt. However, the prompt did not show up.

Is there a way for me to use jax.debug.breakpoint() with Pax? Or do you have an alternative way to debug?

I tried jax.debug.print and it didn't work either.

What pax version are you using?

0.1.0b+2f5c2ed

What jax/jaxlib version are you using?

0.4.20

Which accelerator(s) are you using?

CPU

Additional system info

Python 3.11

NVIDIA GPU info

No response

[Aidandos]

Can you just pass it through the runner state?

[cool-RR]

I'm not sure what you mean. Say that I'm in the _step method of IteratedMatrixGame and I want to know the values of particular variables there, like r2 for example. How might I do the procedure you're suggesting so I'll see its value?