Support Keras 3

I am currently working on using Keras 3 with TF2 backend in RLlib.
But I am facing a few issues.
To fix the issues I created two pull requests. Could anyone please take a look at these and see if we could merge them?

I would really appretierte it if you could look into supporting Keras 3 with different backends.
Currtenlty I can use Keras 3 with legacy mode, but not with the new version.

@arturn could you help me again to get these changes merged?