Running something scan a dynamic number of times

[cool-RR]

Hi everyone! Haven't asked a question here in a while, I hope you missed me 😇

I'm training an actor-critic agent. Every time I train the actor, I want to train the critic as many times as needed until it reaches a good representation of the Q values of the environment. This could be one iterations or 100 iterations. I currenly run these iterations using jax.lax.scan. Does JAX allow the kind of logic I want in a jitted function? i.e. keep running critic training iterations until the loss reaches under a certain condition is satisfied?

[cool-RR]

Oh, I discovered while_loop! That's probably the solution to my problem. Thanks anyway.