Brax fixes

[amrzv]

Hello.

• added missing dependency gymnasium

ModuleNotFoundError                       Traceback (most recent call last)
[<ipython-input-18-e6273b1b5a4b>](https://localhost:8080/#) in <cell line: 2>()
      1 import yaml
----> 2 from rl_games.torch_runner import Runner
      3 
      4 env_name = 'ant'  # @param ['ant', 'humanoid']
      5 configs = {

5 frames
[/usr/local/lib/python3.10/dist-packages/rl_games/torch_runner.py](https://localhost:8080/#) in <module>
      9 from rl_games.common import tr_helpers
     10 
---> 11 from rl_games.algos_torch import a2c_continuous
     12 from rl_games.algos_torch import a2c_discrete
     13 from rl_games.algos_torch import players

[/usr/local/lib/python3.10/dist-packages/rl_games/algos_torch/a2c_continuous.py](https://localhost:8080/#) in <module>
----> 1 from rl_games.common import a2c_common
      2 from rl_games.algos_torch import torch_ext
      3 
      4 from rl_games.algos_torch import central_value
      5 from rl_games.common import common_losses

[/usr/local/lib/python3.10/dist-packages/rl_games/common/a2c_common.py](https://localhost:8080/#) in <module>
      2 import os
      3 
----> 4 from rl_games.common import vecenv
      5 
      6 from rl_games.algos_torch.moving_mean_std import GeneralizedMovingStats

[/usr/local/lib/python3.10/dist-packages/rl_games/common/vecenv.py](https://localhost:8080/#) in <module>
      1 from rl_games.common.ivecenv import IVecEnv
----> 2 from rl_games.common.env_configurations import configurations
      3 from rl_games.common.tr_helpers import dicts_to_dict_with_arrays
      4 import numpy as np
      5 import gym

[/usr/local/lib/python3.10/dist-packages/rl_games/common/env_configurations.py](https://localhost:8080/#) in <module>
      1 import rl_games.envs.test
----> 2 from rl_games.common import wrappers
      3 from rl_games.common import tr_helpers
      4 from rl_games.envs.brax import create_brax_env
      5 from rl_games.envs.envpool import create_envpool

[/usr/local/lib/python3.10/dist-packages/rl_games/common/wrappers.py](https://localhost:8080/#) in <module>
----> 1 import gymnasium
      2 import numpy as np
      3 from numpy.random import randint
      4 
      5 import os

ModuleNotFoundError: No module named 'gymnasium'

and ray with optional = true

ModuleNotFoundError                       Traceback (most recent call last)
[<ipython-input-35-e6273b1b5a4b>](https://localhost:8080/#) in <cell line: 2>()
      1 import yaml
----> 2 from rl_games.torch_runner import Runner
      3 
      4 env_name = 'ant'  # @param ['ant', 'humanoid']
      5 configs = {

4 frames
[/usr/local/lib/python3.10/dist-packages/rl_games/torch_runner.py](https://localhost:8080/#) in <module>
      9 from rl_games.common import tr_helpers
     10 
---> 11 from rl_games.algos_torch import a2c_continuous
     12 from rl_games.algos_torch import a2c_discrete
     13 from rl_games.algos_torch import players

[/usr/local/lib/python3.10/dist-packages/rl_games/algos_torch/a2c_continuous.py](https://localhost:8080/#) in <module>
----> 1 from rl_games.common import a2c_common
      2 from rl_games.algos_torch import torch_ext
      3 
      4 from rl_games.algos_torch import central_value
      5 from rl_games.common import common_losses

[/usr/local/lib/python3.10/dist-packages/rl_games/common/a2c_common.py](https://localhost:8080/#) in <module>
      2 import os
      3 
----> 4 from rl_games.common import vecenv
      5 
      6 from rl_games.algos_torch.moving_mean_std import GeneralizedMovingStats

[/usr/local/lib/python3.10/dist-packages/rl_games/common/vecenv.py](https://localhost:8080/#) in <module>
    114 
    115 
--> 116 class RayVecEnv(IVecEnv):
    117     """Main env class that manages several `rl_games.common.vecenv.Rayworker` objects for parallel training
    118 

[/usr/local/lib/python3.10/dist-packages/rl_games/common/vecenv.py](https://localhost:8080/#) in RayVecEnv()
    121 
    122     """
--> 123     import ray
    124 
    125     def __init__(self, config_name, num_actors, **kwargs):

ModuleNotFoundError: No module named 'ray'

• added brax.v1 into deprecated imports, see Migrate brax v2 in notebook brax_training #302
• cleaned-up brax_training.ipynb notebook
  P.S: on the very last cell the notebook still failing

AttributeError                            Traceback (most recent call last)
[<ipython-input-13-109b243e1a36>](https://localhost:8080/#) in <cell line: 1>()
----> 1 display(visualize(env.env._env.sys, qps))

3 frames
[<ipython-input-12-2430184e7f51>](https://localhost:8080/#) in visualize(sys, qps)
      1 def visualize(sys, qps):
----> 2     return HTML(html.render(sys, qps))

[/usr/local/lib/python3.10/dist-packages/brax/io/html.py](https://localhost:8080/#) in render(sys, states, height, colab, base_url)
     74     string containing HTML for the brax visualizer
     75   """
---> 76   return render_from_json(json.dumps(sys, states), height, colab, base_url)

[/usr/local/lib/python3.10/dist-packages/brax/io/json.py](https://localhost:8080/#) in dumps(sys, states)
    108     RuntimeError: if states have invalid shape
    109   """
--> 110   if any((len(s.x.pos.shape), len(s.x.rot.shape)) != (2, 2) for s in states):
    111     pos_shape = max(len(s.x.pos.shape) for s in states)
    112     rot_shape = max(len(s.x.rot.shape) for s in states)

[/usr/local/lib/python3.10/dist-packages/brax/io/json.py](https://localhost:8080/#) in <genexpr>(.0)
    108     RuntimeError: if states have invalid shape
    109   """
--> 110   if any((len(s.x.pos.shape), len(s.x.rot.shape)) != (2, 2) for s in states):
    111     pos_shape = max(len(s.x.pos.shape) for s in states)
    112     rot_shape = max(len(s.x.rot.shape) for s in states)

AttributeError: 'QP' object has no attribute 'x'

Looks like class QP from cell above needs to be updated to contain x attribute as well apart from pos and rot.

[dmakoviichuk-tt]

Thank you! Ill take a look.