Tutorial 4 — Reinforcement Learning

forge3d ships two Gymnasium-compatible environments for robot learning. Training runs headless (no renderer), and with JAX vmap you can step thousands of parallel environments at near-native speed.

---

ReachEnv — robot reaching

from forge3d.sim.jax_batch import batch_reach_reset, batch_reach_step
import gymnasium as gym

The ReachEnv is available via the forge3d.sim module or as a direct Gymnasium env.

Observation (12 dims)

Index	Description
0–5	Joint angles q₀–q₅ (rad)
6–8	End-effector position (m)
9–11	Target position (m)

Action (6 dims)

Delta joint angles Δq ∈ [−1, 1] rad/step.

---

Training with SB3

import gymnasium as gym
from stable_baselines3 import PPO

# Import the Gymnasium wrapper directly
from forge3d.sim.jax_batch import make_reach_env

env = make_reach_env()         # headless, fast
model = PPO("MlpPolicy", env, verbose=1, n_steps=2048, batch_size=64)
model.learn(total_timesteps=200_000)
model.save("reach_policy")

---

Rendering a rollout

import forge3d as f3d
import numpy as np

world = f3d.World(gravity=(0, 0, -9.81))
world.add_ground()
# … add robot arm …

model = PPO.load("reach_policy")

rec = f3d.Recorder(world, mode="hq", output="rollout.mp4")
rec.run_policy(model, env=None, duration=5.0)

---

JAX batch stepping — 2,000× speedup

import jax
import jax.numpy as jnp
from forge3d.sim.jax_batch import batch_reach_reset, batch_reach_step

key = jax.random.PRNGKey(42)
q, tgt, obs = batch_reach_reset(key, n_envs=256)

# One batched step across 256 environments — single JIT kernel
q, obs, rew, done = batch_reach_step(q, tgt, jnp.zeros((256, 6)))
print(f"rewards: {rew.shape}")  # (256,)

JAX JIT + vmap runs all 256 environments in a single compiled kernel. No Python loop overhead, no device-host round-trips per step.

---

Domain randomization

from forge3d.sim.domain_rand import DomainRand

rand = DomainRand(
    mass_scale=(0.8, 1.2),       # uniform random mass factor
    friction_range=(0.3, 0.9),   # random friction coefficient
    gravity_noise=0.1,           # small perturbation in gravity magnitude
    joint_damping=(0.0, 0.05),   # random damping per joint
)

# Apply randomization to a world at episode start
world = f3d.World()
rand.apply(world, key=jax.random.PRNGKey(episode))

Domain randomization improves sim-to-real transfer by training the policy across a distribution of physical parameters rather than a single point.

---

Switch to JAX backend

ENGINE_BACKEND=jax python train.py

The entire physics stack — including collision detection and contact solving — runs under JAX JIT when ENGINE_BACKEND=jax. This allows jax.grad through the physics step for model-based RL.

---

PickPlaceEnv — pick and place

The pick-and-place environment uses the weld kinematic constraint to simulate grasping.

from forge3d.sim.jax_batch import make_pick_place_env

env = make_pick_place_env()
obs, info = env.reset()
print(f"Obs shape: {obs.shape}")   # (18,)

Observation (18 dims)

Index	Description
0–5	Joint angles q₀–q₅ (rad)
6–8	End-effector position (m)
9–11	Object position (m)
12–14	Object orientation (Euler, rad)
15–17	Target position (m)