Running inference on the model in Realtime

[GandharvMahajan]

Hi, I followed the quadruped joystick controller collab tutorial and trained a PPO policy to walk. I have made some changes at the end for taking user input in realtime (up, down, left and right). I was wondering if it’s possible to run the trained model for realtime inference. Currently I noticed that stepping the jitted model for even 1000 steps take about 20 seconds. Can someone point me to a solution?

thanks.

[erikfrey]

Hi @GandharvMahajan - running a jitted version of the inference function on its own should be ludicrously fast. If you can make a repro in a colab, we can help you investigate.

[GandharvMahajan]

I am sorry, you're right inference is actually really fast. For some reason I thought 3000 steps would correspond to 3 second long "policy" but it is actually a minute long and the jitted functions calculate the policy in about 10 seconds. Thanks. Closing this issue/discussion.

[GandharvMahajan]

Here's the inference code if someone is interested.

`env_fun = create(env_name='barkour', backend='mjx')
ppo_network = ppo.ppo_networks.make_ppo_networks(env_fun.observation_size, env_fun.action_size, preprocess_observations_fn=ppo.running_statistics.normalize, policy_hidden_layer_sizes=(128, 128, 128, 128))
make_inference = ppo.ppo_networks.make_inference_fn(ppo_network)
param_file = 'File_path'
params = model.load_params(param_file)
inference_fn = make_inference(params)
jit_env_reset = jax.jit(env_fun.reset)
jit_env_step = jax.jit(env_fun.step)
jit_inference_fn = jax.jit(inference_fn)

x_vel = 1.0  #@param {type: "number"}
y_vel = 0.0  #@param {type: "number"}
ang_vel = 0.0  #@param {type: "number"}

the_command = jp.array([x_vel, y_vel, ang_vel])

rng = jax.random.PRNGKey(0)
state = jit_env_reset(rng=rng)

state.info['command'] = the_command
rollout = [state.pipeline_state]

for _ in range(3000):
    act, rng = jax.random.split(rng)
    act,_ = jit_inference_fn(state.obs, act)
    state = jit_env_step(state, act)
    rollout.append(state.pipeline_state)
`