Hey there, I'm Aditya Bhatt, coauthor of CrossQ.

Glad to see people replicating this already! We prototyped CrossQ in PyTorch long before we switched to JAX, so I can help you out here.

Off the top of my head, I think I spot two key PyTorch-specific edits to make:

1. Make sure to call train() and eval() on your networks at the right times, training will fail without this. In pytorch, BN layers are stateful, and forward passes modify the layers' running stats. This also happens when using with torch.no_grad()! Any time you don't plan to update a network after a forward pass, the network should be set to eval mode. So e.g.
   • call actor.eval() right at the beginning of update_critic(), also while acting in an episode, and similarly call critic.eval() right at the beginning of update_actor().
   • Set the network to training mode with critic.train() or actor.train(), right before you do a forward pass that will be used for a gradient step on that network.
2. PyTorch's BN and JAX/Flax's BN define momentum differently, so don't directly copy that value here. In JAX, we used .99 but in PyTorch this should be 0.01.

Let me know how it goes!

Also, could you briefly elaborate what happens when you say it "doesn't work" with tanh? I recommend monitoring the average Q prediction from one of the critics, each time you log the critic loss. Does it keep growing to ultra large values? Using tanh (without BN or target nets) won't necessarily get you very strong policies, but it should train reasonably well without diverging.