튜토리얼 4 — 강화학습

forge3d는 로봇 학습을 위한 두 가지 Gymnasium 호환 환경을 제공합니다. 학습은 헤드리스(렌더러 없음)로 실행되며, JAX vmap을 사용하면 거의 네이티브 속도로 수천 개의 병렬 환경을 스텝할 수 있습니다.

---

ReachEnv — 로봇 도달

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

ReachEnv는 forge3d.sim 모듈이나 직접 Gymnasium 환경으로 사용 가능합니다.

관측 (12차원)

인덱스	설명
0–5	조인트 각도 q₀–q₅ (rad)
6–8	엔드이펙터 위치 (m)
9–11	타겟 위치 (m)

행동 (6차원)

조인트 각도 변화량 Δq ∈ [−1, 1] rad/스텝.

---

SB3로 학습

import gymnasium as gym
from stable_baselines3 import PPO
from forge3d.sim.jax_batch import make_reach_env

env = make_reach_env()         # 헤드리스, 빠름
model = PPO("MlpPolicy", env, verbose=1, n_steps=2048, batch_size=64)
model.learn(total_timesteps=200_000)
model.save("reach_policy")

---

롤아웃 렌더링

import forge3d as f3d
import numpy as np

world = f3d.World(gravity=(0, 0, -9.81))
world.add_ground()
# … 로봇 팔 추가 …

model = PPO.load("reach_policy")

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

---

JAX 배치 스텝 — 2,000× 속도 향상

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)

# 256개 환경에 걸친 배치 스텝 — 단일 JIT 커널
q, obs, rew, done = batch_reach_step(q, tgt, jnp.zeros((256, 6)))
print(f"보상 형태: {rew.shape}")  # (256,)

JAX JIT + vmap은 256개 환경 전체를 단일 컴파일 커널로 실행합니다. Python 루프 오버헤드 없음, 스텝당 디바이스-호스트 왕복 없음.

---

도메인 무작위화

from forge3d.sim.domain_rand import DomainRand

rand = DomainRand(
    mass_scale=(0.8, 1.2),       # 균등 분포 질량 인수
    friction_range=(0.3, 0.9),   # 무작위 마찰 계수
    gravity_noise=0.1,           # 중력 크기 소규모 섭동
    joint_damping=(0.0, 0.05),   # 조인트별 무작위 감쇠
)

# 에피소드 시작 시 월드에 무작위화 적용
world = f3d.World()
rand.apply(world, key=jax.random.PRNGKey(episode))

도메인 무작위화는 단일 물리 파라미터 점 대신 분포에 걸쳐 정책을 학습함으로써 시뮬레이션-실제 전이를 향상시킵니다.

---

JAX 백엔드로 전환

ENGINE_BACKEND=jax python train.py

ENGINE_BACKEND=jax로 설정하면 충돌 감지와 접촉 해결을 포함한 전체 물리 스택이 JAX JIT 아래에서 실행됩니다. 이를 통해 모델 기반 RL을 위해 물리 스텝을 통한 jax.grad가 가능합니다.

---

PickPlaceEnv — 집기 & 놓기

집기 및 놓기 환경은 weld 운동학 구속을 사용하여 잡기를 시뮬레이션합니다.

from forge3d.sim.jax_batch import make_pick_place_env

env = make_pick_place_env()
obs, info = env.reset()
print(f"관측 형태: {obs.shape}")   # (18,)

관측 (18차원)

인덱스	설명
0–5	조인트 각도 q₀–q₅ (rad)
6–8	엔드이펙터 위치 (m)
9–11	물체 위치 (m)
12–14	물체 방향 (오일러, rad)
15–17	타겟 위치 (m)