[TEST] Jax Separate Networks Loss

[RuanJohn]

What do you want to test?

Unit test for the MAPGWithTrustRegionClippingLossSeparateNetworks component of the Jax PPO implementation that makes use of separate critic and policy netowrks.

Outline of test structure

• Unit tests
• Test components and hooks

Definition of done

Passing checks, cover all hooks, edge cases considered

Mandatory checklist before making a PR

• [ ] The success criteria laid down in “Definition of done” are met.
• [ ] Test code is documented - docstrings for methods and classes, static types for arguments.
• [ ] Documentation is updated - README, CONTRIBUTING, or other documentation.