Return optional MJX substep states

[sylvesterkaczmarek]

Summary

Adds an opt-in return_substeps argument to mjx_env.step so callers can accumulate rewards or metrics over frame-skipped simulation states.

The existing path is unchanged by default and continues to return only the final mjx.Data, without storing scan outputs. With return_substeps=True, the function returns (final_data, substeps), where the substep pytree has a leading axis of n_substeps.

Addresses #296.

Tests

• Added a focused MJX regression covering three returned substeps.
• Verified substep times, final-state equality, and backward-compatible default output.
• pytest mujoco_playground/_src/mjx_env_test.py: passed.
• Ruff and Pyink checks for the new test file: passed.