How to render test results of trained policy?

[hdadong]

@jake-levy I ran your go1_render.py, but I found the following bug:

(ssrl) (base) bigeast@bigeast:~/ssrl/ssrl$ python scripts/go1_render.py
/home/bigeast/ssrl/ssrl/scripts/go1_render.py:12: UserWarning: 
The version_base parameter is not specified.
Please specify a compatability version level, or None.
Will assume defaults for version 1.1
  @hydra.main(config_path="configs", config_name="go1")
/home/bigeast/anaconda3/envs/ssrl/lib/python3.9/site-packages/hydra/_internal/hydra.py:119: UserWarning: Future Hydra versions will no longer change working directory at job runtime by default.
See https://hydra.cc/docs/1.2/upgrades/1.1_to_1.2/changes_to_job_working_dir/ for more information.
  ret = run_job(
[2024-10-17 19:30:08,364][jax._src.xla_bridge][INFO] - Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
[2024-10-17 19:30:08,364][jax._src.xla_bridge][INFO] - Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: libtpu.so: cannot open shared object file: No such file or directory
2024-10-17 19:30:08.377632: W external/xla/xla/service/gpu/nvptx_compiler.cc:765] The NVIDIA driver's CUDA version is 12.2 which is older than the ptxas CUDA version (12.6.77). Because the driver is older than the ptxas version, XLA is disabling parallel compilation, which may slow down compilation. You should update your NVIDIA driver or use the NVIDIA-provided CUDA forward compatibility packages.
Error executing job with overrides: []
Traceback (most recent call last):
  File "/home/bigeast/ssrl/ssrl/scripts/go1_render.py", line 15, in render_go1
    from brax.envs.go1_deltaxy_pd_slow import Go1DeltaXyPdSlow as Go1
ModuleNotFoundError: No module named 'brax.envs.go1_deltaxy_pd_slow'

Set the environment variable HYDRA_FULL_ERROR=1 for a complete stack trace.

Then I modified go1_render.py, the code is as follows:

"""Run this script with:
python -m streamlit run go1_render.py"""

import dill
import hydra
from omegaconf import DictConfig
from pathlib import Path
import functools as ft
import streamlit as st

@hydra.main(config_path="configs", config_name="go1")
def render_go1(cfg: DictConfig):

    from brax.envs.go1_go_fast import Go1GoFast
    from brax.io.rendering import render

    # load states to render
    render_name = "ssrl_ssrl_2024-10-17_1538_32_step_14000_rew_1122.pkl"#f"go1_{cfg.render.policy}_render.pkl"
    path = (Path(__file__).parent.parent
            / 'saved_policies'
            / render_name)
    with open(path, 'rb') as f:
        pipeline_states = dill.load(f)
    env_kwargs = cfg.env_ssrl

    # create env
    env_dict = {'Go1GoFast': Go1GoFast}
    env_fn = ft.partial(env_dict[cfg.env], backend='generalized')
    env_fn = add_kwargs_to_fn(env_fn, **env_kwargs)
    env = env_fn()

    render_heght = 800
    st.set_page_config(layout="wide")
    render(env, pipeline_states, render_heght)

def add_kwargs_to_fn(partial_fn, **kwargs):
    """add the kwargs to the passed in partial function"""
    for param in kwargs:
        partial_fn.keywords[param] = kwargs[param]
    return partial_fn

if __name__ == '__main__':
    render_go1()

But still error:

(ssrl) (base) bigeast@bigeast:~/ssrl/ssrl$ python scripts/go1_render.py
/home/bigeast/ssrl/ssrl/scripts/go1_render.py:12: UserWarning: 
The version_base parameter is not specified.
Please specify a compatability version level, or None.
Will assume defaults for version 1.1
  @hydra.main(config_path="configs", config_name="go1")
/home/bigeast/anaconda3/envs/ssrl/lib/python3.9/site-packages/hydra/_internal/hydra.py:119: UserWarning: Future Hydra versions will no longer change working directory at job runtime by default.
See https://hydra.cc/docs/1.2/upgrades/1.1_to_1.2/changes_to_job_working_dir/ for more information.
  ret = run_job(
[2024-10-17 19:35:02,824][jax._src.xla_bridge][INFO] - Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
[2024-10-17 19:35:02,825][jax._src.xla_bridge][INFO] - Unable to initialize backend 'tpu': INTERNAL: Failed to open libtpu.so: libtpu.so: cannot open shared object file: No such file or directory
2024-10-17 19:35:02.838643: W external/xla/xla/service/gpu/nvptx_compiler.cc:765] The NVIDIA driver's CUDA version is 12.2 which is older than the ptxas CUDA version (12.6.77). Because the driver is older than the ptxas version, XLA is disabling parallel compilation, which may slow down compilation. You should update your NVIDIA driver or use the NVIDIA-provided CUDA forward compatibility packages.
[2024-10-17 19:35:03,268][root][INFO] - Converting mesh (3718290260903023517, -870898390182723615) into convex hull.
[2024-10-17 19:35:05,858][root][INFO] - Converting mesh (4320480779867028070, -4780402356330154556) into convex hull.
[2024-10-17 19:35:07,199][root][INFO] - Converting mesh (3170428720903045684, -7216126722741586661) into convex hull.
[2024-10-17 19:35:09,263][root][INFO] - Converting mesh (-2877919857730206175, 3312743317921156319) into convex hull.
[2024-10-17 19:35:11,283][root][INFO] - Converting mesh (4407398090165108200, -1590417479951165338) into convex hull.
2024-10-17 19:35:14.274 WARNING streamlit.runtime.scriptrunner_utils.script_run_context: Thread 'MainThread': missing ScriptRunContext! This warning can be ignored when running in bare mode.
Error executing job with overrides: []
Traceback (most recent call last):
  File "/home/bigeast/ssrl/ssrl/scripts/go1_render.py", line 35, in render_go1
    render(env, pipeline_states, render_heght)
  File "/home/bigeast/ssrl/ssrl/brax/io/rendering.py", line 14, in render
    render = html.render(env.sys.replace(dt=env.dt), pipeline_states,
  File "/home/bigeast/ssrl/ssrl/brax/io/html.py", line 73, in render
    return render_from_json(json.dumps(sys, states), height, colab, base_url)
  File "/home/bigeast/ssrl/ssrl/brax/io/json.py", line 104, in dumps
    if any((len(s.x.pos.shape), len(s.x.rot.shape)) != (2, 2) for s in states):
  File "/home/bigeast/ssrl/ssrl/brax/io/json.py", line 104, in <genexpr>
    if any((len(s.x.pos.shape), len(s.x.rot.shape)) != (2, 2) for s in states):
AttributeError: 'RunningStatisticsState' object has no attribute 'x'

Set the environment variable HYDRA_FULL_ERROR=1 for a complete stack trace.

Could you please provide the complete testing and rendering code? I believe this will encourage people to follow your awesome work.

[jake-levy]

I added rendering instructions to the simulation readme.

[hdadong]

How to save the world model in the training? And evaluate the world model and policy after training?

[jake-levy]

When you say "evaluate the world model", do you mean how to generate Figure 5 (right)?

[hdadong]

Since my goal is to train a humanoid robot, direct training in the real world is not feasible. One approach with SSRL is to pretrain the world model and policy in simulation, then continue training in the real world. If you have any code available for testing the world model and policy after simulation training, that would be very helpful.

[jake-levy]

I see. I can throw together something to eval and render a saved policy. Reopening...

[hdadong]

Thanks. Sorry to confuse, when I say "evaluate the world model", I mean evaluate the torque estimator network.

[jake-levy]

Right. So for evaluating the torque estimator network, are you referring to comparing torque outputs from the network against actual torque values from the simulator?

[hdadong]

Yeah!