Why Jumpy?

[cemlyn007]

Wondering what was the reason for creating Brax to allow NumPy and JAX via jumpy?

Why not just use JAX completely otherwise you have to add a method to jumpy that exists in NumPy and JAX? Then things would be simpler and easier to test (for me at least)?

Maybe I am the only one missing something here. 😅

[cemlyn007]

For example what should be the expected outcome of:

import brax
import brax.jumpy as jp
env = brax.envs.create('ant')
state = env.reset(jp.random_prngkey(0))
state = env.step(state, jp.random_uniform(jp.random_prngkey(0), (env.action_size,), -1., 1.))

When using Brax 0.0.15. Should I expect an error?

[erikfrey]

Jumpy allows us to operate Brax in a sort of "immediate mode", where things like physics steps can be run without any jitting overhead - this lowers the time to run a single physics step from tens of seconds to less than a second. This is very useful when debugging environments or the underlying engine, where you want to make lots of small changes and review each change quickly.

Since we wrote Jumpy, JAX has improved in its ability to be run without jax.jit and return quickly. So one day we may review this decision and it's time to retire Jumpy... but not today.

[cemlyn007]

Thank you for your detailed response. Have a good weekend!