https://github.com/flairox/jaxgl
Simple JAX Graphics Library.
https://github.com/flairox/jaxgl
graphics graphics-library jax
Last synced: 4 months ago
JSON representation
Simple JAX Graphics Library.
- Host: GitHub
- URL: https://github.com/flairox/jaxgl
- Owner: FLAIROx
- License: mit
- Created: 2024-10-30T20:33:39.000Z (about 1 year ago)
- Default Branch: main
- Last Pushed: 2024-11-03T22:07:04.000Z (about 1 year ago)
- Last Synced: 2025-09-17T07:21:08.871Z (4 months ago)
- Topics: graphics, graphics-library, jax
- Language: Python
- Homepage: https://kinetix-env.github.io/
- Size: 60.5 KB
- Stars: 36
- Watchers: 4
- Forks: 1
- Open Issues: 0
-
Metadata Files:
- Readme: README.md
- License: LICENSE
Awesome Lists containing this project
README
# JaxGL
JaxGL is a simple and flexible graphics library written entirely in JAX. JaxGL was created by [Michael Matthews](https://github.com/MichaelTMatthews) and [Michael Beukman](https://github.com/Michael-Beukman) for the [Kinetix](https://github.com/FLAIROx/Kinetix) project.
# 💻 Basic Usage
```python
# 512x512 pixels
screen_size = (512, 512)
# Clear a fresh screen with a black background
clear_colour = jnp.array([0.0, 0.0, 0.0])
pixels = clear_screen(screen_size, clear_colour)
# We render to a 256x256 'patch'
patch_size = (256, 256)
triangle_renderer = make_renderer(screen_size, fragment_shader_triangle, patch_size)
# Patch position (top left corner)
pos = jnp.array([128, 128])
triangle_data = (
# Vertices (note these must be anti-clockwise)
jnp.array([[150, 200], [150, 300], [300, 150]]),
# Colour
jnp.array([255.0, 0.0, 0.0]),
)
# Render the triangle to the screen
pixels = triangle_renderer(pixels, pos, triangle_data)
```
This produces the following image:
# 👨💻 Custom Shaders
Arbitrary rendering effects can be achieved by writing your own shaders.
```python
screen_size = (512, 512)
clear_colour = jnp.array([0.0, 0.0, 0.0])
pixels = clear_screen(screen_size, clear_colour)
patch_size = (256, 256)
# We make our own variation of the circle shader
# We give both a central and edge colour and interpolate between these
# Each fragment shader has access to
# position: global position in screen space
# current_frag: the current colour of the fragment (useful for transparency)
# unit_position: the position inside the patch (scaled to between 0 and 1)
# uniform: anything you want for your shader. These are the same for every fragment.
def my_shader(position, current_frag, unit_position, uniform):
centre, radius, colour_centre, colour_outer = uniform
dist = jnp.sqrt(jnp.square(position - centre).sum())
colour_interp = dist / radius
colour = colour_interp * colour_outer + (1 - colour_interp) * colour_centre
return jax.lax.select(dist < radius, colour, current_frag)
circle_renderer = make_renderer(screen_size, my_shader, patch_size)
# Patch position (top left corner)
pos = jnp.array([128, 128])
# This is the uniform that is passed to the shader
circle_data = (
# Centre
jnp.array([256.0, 256.0]),
# Radius
100.0,
# Colour centre
jnp.array([255.0, 0.0, 0.0]),
# Colour outer
jnp.array([0.0, 255.0, 0.0]),
)
# Render the triangle to the screen
pixels = circle_renderer(pixels, pos, circle_data)
```
# 🔄 In Kinetix
JaxGL is used for rendering in [Kinetix](https://github.com/FLAIROx/Kinetix). Shown below is an example robotics grasping task.
# ⬇️ Installation
To use JaxGL in your work you can install via PyPi:
```commandline
pip install jaxgl
```
If you want to extend JaxGL you can install as follows:
```commandline
git clone https://github.com/FLAIROx/JaxGL
cd JaxGL
pip install -e ".[dev]"
pre-commit install
```
# 🔍 See Also
- [JAX Renderer](https://github.com/JoeyTeng/jaxrenderer) A more complete JAX renderer more suitable for 3D rendering.
- [Jax2D](https://github.com/MichaelTMatthews/Jax2D) 2D physics engine in JAX.
- [Kinetix](https://github.com/FLAIROx/Kinetix) physics-based reinforcement learning in JAX.