Add real-time rendering callback #634
Conversation
Modified for Viewer
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
| return | ||
| io_callback(render_fn, None, state) | ||
|
|
||
| jax.lax.cond(should_render, render, lambda s: None, nstate) |
There was a problem hiding this comment.
can this whole block just be
if render_fn:
io_callback(render_fn, None, state)
There was a problem hiding this comment.
This may get rid of the fixed overhead in your main post, JAX should be ignoring this whole block
| eval_policy_fn(policy_params), | ||
| key, | ||
| unroll_length=episode_length // action_repeat, | ||
| should_render=jnp.array(False, dtype=jnp.bool_), # No rendering during eval |
There was a problem hiding this comment.
you're not passing render_fn anyways, not sure you really need should_render
| policy_params_fn: Callable[..., None] = lambda *args: None, | ||
| # rendering | ||
| render_fn: Optional[Callable[[envs.State], None]] = None, | ||
| should_render: jax.Array = jnp.array(True, dtype=jnp.bool_), |
| Returns: | ||
| Tuple of (make_policy function, network params, metrics) | ||
| """ | ||
| # If the environment is wrapped with ViewerWrapper, use its rendering functions. |
There was a problem hiding this comment.
not sure what ViewerWrapper is, maybe update the comment?
| unroll_length, | ||
| extra_fields=('truncation', 'episode_metrics', 'episode_done'), | ||
| render_fn=render_fn, | ||
| should_render=should_render, |
There was a problem hiding this comment.
again, you maybe can get away without the bool should_render
| # optimization | ||
|
|
||
| # check for rendering dynamically | ||
| should_render_py = False |
There was a problem hiding this comment.
so you're ignoring the arg to train(), just delete all the args
This PR adds an interface for real-time rendering in
PPO.trainandactingby adding a user-defined callback executed viajax.experimental.io_callback. The final goal is to provide a real-time viewer beside the notebook html viewer.Key Changes
Two new optional parameters are added:
render_fn: A Python callable that accepts abrax.Stateto handle the rendering logic.should_render: A boolean JAXArrayused to conditionally trigger the callback.Performance Impact
Adding the new callback introduces a minor performance. This overhead exists even when rendering is disabled (i.e.,
should_renderisFalse). The JIT compiler must account for the conditional logic required for theio_callback, which slightly alters the compiled execution path.The benchmarks below, run on an Apple M1 Max and NVIDIA 2080.