Using Jax as learning framework

[jrabary]

Is it possible to use Jax as learning framework without any performance/speed of simulation loss ?

[Mayankm96]

Natively, Isaac Sim supports numpy and torch backend, while Orbit right now mainly uses torch since that is more useful for GPU parallelization. Additionally, for small-dimension vectors, there isn't a lot of overhead moving them from torch to numpy when needed (such as for ROS). Our future plan is to support warp backend since that would simplify the processing of large image observations in a batched manner.

We haven't yet investigated using Jax. From what I saw here, there isn't any Pytorch<->Jax operations available right now. It would definitely be useful for researchers since Jax has its benefits.

Are there any Gym environments that are supporting Jax? We can take a look and see how much effort this would take.

[StoneT2000]

Some gym envs in jax: https://github.com/RobertTLange/gymnax
Then there's also Brax

For converting jax to torch tensors directly on cuda, you can use dlpack as so

import jax
import jax.dlpack
import torch
import torch.utils.dlpack

def jax_to_torch(x):
    return torch.utils.dlpack.from_dlpack(jax.dlpack.to_dlpack(x))
def torch_to_jax(x):
    return jax.dlpack.from_dlpack(torch.utils.dlpack.to_dlpack(x))

a = torch.tensor([1,2,3]).cuda()
a_jax = torch_to_jax(a)
print(a_jax)
# both a and a_jax are on GPU

[Mayankm96]

Nice! Thanks a lot for pointing this out @StoneT2000

Is there significant overhead in these conversions when dealing with large tensors? If not, then I am happy to have an interface that converts the environments to be compatible with JAX-based libraries.

[StoneT2000]

Maybe? I myself haven't tested it just yet at scale altho I plan to soon, would be nice to see a speed test

[Mayankm96]

I tried to make a script out of the code mentioned above but haven't been able to run it. Always get some ModuleImportError when loading jax. Currently, there are other tasks queued up for us which are on a higher priority.

I am leaving the script here for someone to try it out :)

import torch
import torch.utils.dlpack
import jax
import jax.dlpack

import time

# A generic mechanism for turning a JAX function into a PyTorch function.

def j2t(x_jax):
  x_torch = torch.utils.dlpack.from_dlpack(jax.dlpack.to_dlpack(x_jax))
  return x_torch

def t2j(x_torch):
  x_torch = x_torch.contiguous()  # https://github.com/google/jax/issues/8082
  x_jax = jax.dlpack.from_dlpack(torch.utils.dlpack.to_dlpack(x_torch))
  return x_jax

# time the JAX version
x = torch.randn(2048, 512, 4).cuda()
start = time.perf_counter()
for _ in range(100):
  y = t2j(x)
print(f"JAX time    : {(time.perf_counter() - start) / 100:.6f} s")

# time the PyTorch version
x = jax.random.normal(jax.random.PRNGKey(0), (2048, 512, 4))
start = time.perf_counter()
for _ in range(100):
  y = j2t(x)
print(f"PyTorch time: {(time.perf_counter() - start) / 100:.6f} s")

[StoneT2000]

Modified the code to run N=100_000 times for each test

JAX time    : 0.032 ms
PyTorch time: 0.003 ms

Seems like no serious performance drop, and obviously transferring larger matrices in fewer loops it makes it even faster

[Mayankm96]

I fixed the script a bit. The perf_counter reports time in seconds. So that makes it 32 ms and 3 ms. Is this the time taken for N=100_000? If that's the case, it is quite small. We can add this functionality as a wrapper.

[StoneT2000]

Yes 32ms and 3ms are the time for 100_000 so this is more than fast enough if each individual op is taking 32 / 100_000 ms. Would be great to use a wrapper! I was unaware about the contiguous issue before as well so good catch.

[Mayankm96]

Since we are running behind our current development roadmap, it is a bit difficult to allocate immediate resources onthis wrapper.

@StoneT2000 would this be something you'd be interested and have time to contribute into the framework? :)

[StoneT2000]

Sure thing, I already have one I use myself. I assume it would go somewhere in the utils/wrappers folder?

It seems existing wrappers are on a library by library basis, would it still go in that folder?

If so I can port over and clean up a version that just converts outputs to jax arrays (jax library specific codes are a future thing I guess, the most popular ones I'm aware of are acme and jaxrl)

[Mayankm96]

Hi @StoneT2000 , sorry for the late reply.

Yes, the wrapper can go into the utils/wrappers directory. You can name it something like jax.py if it makes sense.

[Mayankm96]

For people further interested in this functionality, I found another library that seems useful: https://github.com/luchris429/purejaxrl

[Toni-SM]

Starting with skrl v1.0.0-rc.1, there is JAX support for Isaac Gym, Omniverse Isaac Gym and Isaac Orbit among others.
Visit the example section for Isaac Orbit in skrl's docs to start training using PyTorch or JAX.

Working on the integration of skrl's JAX examples into the Isaac Orbit extensions :)