updates for brax_visualization.ipynb
Closed this issue · 9 comments
It seems brax changed their API and the brax_visualization.ipynb needs some quick fix:
in cell#3, line#5:
config = runner.get_prebuilt_config()needs to be commented/removed
in cell#3, line#8:
env_config = config['env_config']should change to
env_config = runner.params['config']['env_config']
In cell#5, line#14:
env.state.qp should change to env.env._state.qp
In cell#5, line#17:
env.step(act.unsqueeze(0)) should change to env.step(act)
in cell#7,
display(visualize(env.env.sys, qps))should change to
display(visualize(env.env._env.sys, qps))Thank you,
Ill test it. didn't touch it since their release :)
Even after these changes, I still can't run it with the latest Brax:
ValueError: vmap got inconsistent sizes for array axes to be mapped:
the tree of axis sizes is:
(State(qp=QP(pos=1, rot=1, vel=1, ang=1), obs=1, reward=1, done=1, metrics={'reward_contact_cost': 1, 'reward_ctrl_cost': 1, 'reward_forward': 1, 'reward_survive': 1}, info={'first_obs': 1, 'first_qp': QP(pos=1, rot=1, vel=1, ang=1), 'steps': 1, 'truncation': 1}), 8)
@boxiXia what version of Brax did you use? Or maybe you have some other changes in place?
I just did pip install brax, brax version is: 0.0.10.
The example I tested is documented here
I did pip install from the local folder with Brax repo, after pulling the latest. It was also 0.0.10. But a cell with qps.append(QP(env.env._state.qp)) still fails for me.
Here is the brax_visualization.ipynb jupyter notebook modified for the visualization to work:
https://gist.github.com/boxiXia/1acd8897352a3404a172ee75d984fd20
Still seeing this error. What version of rl_games do you use, from github the latest master or pip install rl-games?