Policy Gradient algorithms' issues

If you compared PG verisons in `jax` and `pytorch` versions of the algorithms, you would notice following issues:
1. `torch` implements its own losses in a separate file, `jax` takes them from `rlax`, the versions do not align
2. `torch` has sample-based `NeuRD` implemented, whereas `jax` — not.
3. `torch` version does have a notion of `batch_size`, `jax`'s version works with a whole buffer
4. there is so much boilerplate repetitions and tweaks in `jax` code in operations like sampling buffer sampling and loss function definitions, there is no native legal action masking support in the loss, only — when acting
5. ppo agent is implemented only in `torch` and for `atari` which is not very useful and might be deprecated

I regard these issues essential, but what's the maintainers priority on resolving them, considering the 2.0 release?
cc @lanctot

> [!IMPORTANT]
> This issue is created to discuss potential changes and *does not imply* that you need to rush a create a PR.