Comments (5)
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
from isaaclab.
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.
from isaaclab.
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.
from isaaclab.
Maybe? I myself haven't tested it just yet at scale altho I plan to soon, would be nice to see a speed test
from isaaclab.
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:.3f} ms")
# 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:.3f} ms")
from isaaclab.
Related Issues (20)
- [Question] How to get the local pose or world pose of an object in IsaacLab? HOT 8
- [Proposal] Config action order HOT 1
- [Bug Report] import omni.replicator.core as rep ModuleNotFoundError: No module named 'omni.replicator' HOT 2
- [Question] train from checkpoint HOT 3
- [Proposal] Adding Semantic Capabilities
- [Bug Report] ValueError: The passed ArgParser object already has the field 'device'. HOT 1
- [Bug Report] libhdx.so: cannot open shared object file: No such file or directory HOT 1
- [Proposal] Update Se3Gamepad to support new gamepad controller HOT 2
- Failed to resolve extension dependencies for isaaclab.python-1.1.0 HOT 1
- [Question] Question about changing joint's drive type HOT 5
- [Bug Report] 'RigidBodyView' object has no attribute 'get_accelerations' HOT 2
- [Bug Report] Wrong SpaceMouse z-axis rotation direction in Isaac-Lift-Cube-Franka-IK-Rel-v0 teleoperation demo
- How to use RayCasterCamera HOT 1
- [Question] Contact force for the soft bodies and gripper HOT 2
- [Bug Report] Adding call for CV2 imshow will freeze entire Isaac-Sim HOT 2
- [Question] Why I can get semantic or instance image for shape cube, but not for franka pandar arm? HOT 3
- [Question] Runtime error with Creating an empty scene HOT 2
- [Question] Distortion coefficient of cameras HOT 1
- [Bug Report] IsaacLab throws exceptions when adding elements from IsaacSim menu
- Updating the camera view using set_world_pose_from_view HOT 3
Recommend Projects
-
React
A declarative, efficient, and flexible JavaScript library for building user interfaces.
-
Vue.js
🖖 Vue.js is a progressive, incrementally-adoptable JavaScript framework for building UI on the web.
-
Typescript
TypeScript is a superset of JavaScript that compiles to clean JavaScript output.
-
TensorFlow
An Open Source Machine Learning Framework for Everyone
-
Django
The Web framework for perfectionists with deadlines.
-
Laravel
A PHP framework for web artisans
-
D3
Bring data to life with SVG, Canvas and HTML. 📊📈🎉
-
Recommend Topics
-
javascript
JavaScript (JS) is a lightweight interpreted programming language with first-class functions.
-
web
Some thing interesting about web. New door for the world.
-
server
A server is a program made to process requests and deliver data to clients.
-
Machine learning
Machine learning is a way of modeling and interpreting data that allows a piece of software to respond intelligently.
-
Visualization
Some thing interesting about visualization, use data art
-
Game
Some thing interesting about game, make everyone happy.
Recommend Org
-
Facebook
We are working to build community through open source technology. NB: members must have two-factor auth.
-
Microsoft
Open source projects and samples from Microsoft.
-
Google
Google ❤️ Open Source for everyone.
-
Alibaba
Alibaba Open Source for everyone
-
D3
Data-Driven Documents codes.
-
Tencent
China tencent open source team.
from isaaclab.