Physics Simulation과 JAX 생태계

JAX가 특히 빛나는 또 하나의 영역은 미분 가능한 물리 시뮬레이션이야. 실제 물리 법칙으로 도는 시뮬레이터를 통해 그래디언트가 흘러서, 시스템 식별, 제어, 강화학습의 새 패러다임이 가능해.

Brax, Google의 JAX 물리 엔진

pip install brax

from brax import envs
from brax.io import html

# pendulum 환경
env = envs.create("pendulum")

# rollout
state = env.reset(jax.random.PRNGKey(0))
for _ in range(200):
    action = jnp.array([0.0])  # zero action
    state = env.step(state, action)

Brax는 JAX 위에서 실행되는 강체 물리 엔진이야. MuJoCo와 비슷한 환경을 다수 제공해. GPU에서 1000개 환경을 vmap으로 동시에 시뮬레이션할 수 있어. 강화학습 속도를 최대 1000배까지 높일 수 있어.

# 1000 개 environment, 동시 rollout
envs_batch = jax.vmap(env.step, in_axes=(0, 0))
states = jax.vmap(env.reset)(jax.random.split(key, 1000))

for _ in range(200):
    actions = policy(states)   # (1000, action_dim)
    states = envs_batch(states, actions)

MJX, MuJoCo on JAX

DeepMind가 2024년에 출시한 MuJoCo의 JAX 포팅이야. Brax보다 접촉과 마찰을 더 정확히 모델링하고 강화학습과 로보틱스의 표준 환경을 다수 제공해.

import mujoco
from mujoco import mjx

model = mujoco.MjModel.from_xml_path("humanoid.xml")
mjx_model = mjx.put_model(model)

@jax.jit
def step(data, action):
    data = data.replace(ctrl=action)
    return mjx.step(mjx_model, data)

# 1000 개 humanoid simulation 동시
batched_step = jax.vmap(step)
batch_data = jax.vmap(mjx.make_data)(...)

JAX-MD, molecular dynamics

pip install jax-md

from jax_md import space, energy, simulate

# 입자 시뮬레이션
displacement_fn, shift_fn = space.periodic(box_size)
energy_fn = energy.lennard_jones_pair(displacement_fn)

init, apply = simulate.nve(energy_fn, shift_fn, dt=1e-3)
state = init(rng, R, mass=1.0)

# rollout
for _ in range(1000):
    state = apply(state)

# 미분 가능 — energy 의 parameter 학습 가능
def loss(epsilon, state):
    energy_fn = energy.lennard_jones_pair(displacement_fn, epsilon=epsilon)
    new_state = apply_with_energy(energy_fn, state)
    return some_objective(new_state)

grad_eps = jax.grad(loss)(epsilon, state)

왜 JAX가 다른가?

항목	전통 (C++/Python)	JAX-based
속도 (CPU)	빠름 (C++)	비슷
속도 (GPU)	드뭄, 별도 코드	자동, 코드 그대로
대량 벡터화	수동 (한 sim 한 thread)	vmap 자동
미분	없거나 hand-coded adjoint	자동
RL 학습 속도	baseline	10-1000x

강화학습과 로보틱스에서는 환경 시뮬레이션이 병목이야. JAX 기반 환경은 시뮬레이션을 크게 가속하고 에이전트 학습도 같은 장치에서 자연스럽게 이어 가.

🌐 산업 / 학술 현황

DeepMind와 Google Research를 비롯한 여러 연구 조직이 새로운 강화학습 코드에 JAX 기반 시뮬레이터를 적극 활용하고 있어. 그 영향은 학계와 스타트업으로 확산 중이야. 로보틱스와 강화학습에서 JAX의 비중은 더 커질 가능성이 높고, 지금은 이를 일찍 익힐 좋은 시기야.

전통 시뮬레이터의 정확성이 더 높은 영역도 있어 (복잡한 접촉이나 연체 동역학 등). JAX 시뮬레이터들은 빠르게 따라잡는 중. 큰 배치와 빠른 학습, 높은 정확성 사이에는 절충이 있어. 강화학습의 많은 사용 사례에서는 큰 배치의 이점이 특히 중요해.

Code

# JAX-MD provides differentiable molecular dynamics simulations
# Example concept (simplified):
import jax
import jax.numpy as jnp

def lennard_jones(r, epsilon=1.0, sigma=1.0):
    """Lennard-Jones potential between two particles at distance r."""
    s6 = (sigma / r) ** 6
    return 4.0 * epsilon * (s6**2 - s6)

# The potential is differentiable — get forces automatically
force = -jax.grad(lennard_jones)
print(f"Force at r=1.5: {force(1.5):.4f}")

# Vectorize over all pairs
def total_energy(positions):
    """Compute total LJ energy for a system of particles."""
    n = positions.shape[0]
    energy = 0.0
    for i in range(n):
        for j in range(i + 1, n):
            r = jnp.linalg.norm(positions[i] - positions[j])
            energy += lennard_jones(r)
    return energy

# Get forces on all particles: F = -∇E
forces = -jax.grad(total_energy)(positions)

# Brax enables training robot controllers through physics
# The entire simulation is differentiable:
# environment_step → reward → grad(policy_params)
# This is vastly faster than finite-difference methods used in
# traditional RL for continuous control.

# JAXopt provides scipy-style optimization with JAX
import jaxopt

# Minimize a function with L-BFGS
def rosenbrock(x):
    return (1 - x[0])**2 + 100 * (x[1] - x[0]**2)**2

solver = jaxopt.LBFGS(fun=rosenbrock, maxiter=100)
result = solver.run(jnp.array([0.0, 0.0]))
print(f"Minimum at: {result.params}")  # ≈ [1.0, 1.0]

# Projected gradient descent (constrained optimization)
solver = jaxopt.ProjectedGradient(
    fun=objective,
    projection=jaxopt.projection.projection_box,
    maxiter=200,
)

import jax
import jax.numpy as jnp

# Compute Taylor coefficients of a function
# This is useful for solving PDEs with high-order derivatives
def f(x):
    return jnp.sin(x) * jnp.exp(-x)

# Get the 4th derivative efficiently
from jax import jet

# jet computes Taylor-mode derivatives
primals = jnp.array(1.0)
series = (jnp.array(1.0), jnp.array(0.0), jnp.array(0.0), jnp.array(0.0))
primal_out, series_out = jax.jet(f, (primals,), (series,))
# series_out contains higher-order derivatives
# This scales linearly with derivative order, not exponentially

🌐 산업 / 학술 현황

Code

External links

Exercise

Progress

댓글 0