Bug Fix & Distributed Training Improvement

[ZaberKo]

Description

This PR includes these modifications:

• Fix bug in brax.
• orbax support for state save&load.
• Unified field method in dataclass.
• State usage improvement.
• New Distributed API

Backward Compatibility Break:

• workflow.step(state) -> (train_info, state)
• New state|stateful dataclass API with pytree_field()

Checklist

• I have formatted my Python code with black.
• I have good commit messages.
• If adding new algorithms, problems, operators:
  • Added related test cases.
  • Added docstring to explain important parameters.
  • Added entries in the docs.