From 52ccf88fbc1c31bbfff6845315dc5dc0ece53b31 Mon Sep 17 00:00:00 2001 From: vyeoms Date: Fri, 21 Aug 2026 22:25:38 +0100 Subject: [PATCH 1/5] Add MuJoCo playground env replicas: Ant, HalfCheetah, Hopper, Humanoid, Walker2D --- .gitignore | 1 + build.sh | 4 + config/mjc_ant.ini | 48 + config/mjc_half_cheetah.ini | 46 + config/mjc_hopper.ini | 47 + config/mjc_humanoid.ini | 48 + config/mjc_walker2d.ini | 47 + ocean/mujoco/backend.h | 129 +++ ocean/mujoco/mjc_ant.c | 3 + ocean/mujoco/mjc_ant.cu | 4 + ocean/mujoco/mjc_ant.h | 176 ++++ ocean/mujoco/mjc_half_cheetah.c | 3 + ocean/mujoco/mjc_half_cheetah.cu | 4 + ocean/mujoco/mjc_half_cheetah.h | 158 +++ ocean/mujoco/mjc_hopper.c | 3 + ocean/mujoco/mjc_hopper.cu | 4 + ocean/mujoco/mjc_hopper.h | 165 ++++ ocean/mujoco/mjc_humanoid.c | 3 + ocean/mujoco/mjc_humanoid.cu | 4 + ocean/mujoco/mjc_humanoid.h | 194 ++++ ocean/mujoco/mjc_walker2d.c | 3 + ocean/mujoco/mjc_walker2d.cu | 4 + ocean/mujoco/mjc_walker2d.h | 160 ++++ ocean/mujoco/mjcf2bin.py | 73 ++ ocean/mujoco/physics.h | 1484 +++++++++++++++++++++++++++++ ocean/mujoco/render.h | 69 ++ resources/mujoco/ant.xml | 81 ++ resources/mujoco/half_cheetah.xml | 96 ++ resources/mujoco/hopper.xml | 53 ++ resources/mujoco/humanoid.xml | 121 +++ resources/mujoco/walker2d.xml | 68 ++ 31 files changed, 3303 insertions(+) create mode 100644 config/mjc_ant.ini create mode 100644 config/mjc_half_cheetah.ini create mode 100644 config/mjc_hopper.ini create mode 100644 config/mjc_humanoid.ini create mode 100644 config/mjc_walker2d.ini create mode 100644 ocean/mujoco/backend.h create mode 100644 ocean/mujoco/mjc_ant.c create mode 100644 ocean/mujoco/mjc_ant.cu create mode 100644 ocean/mujoco/mjc_ant.h create mode 100644 ocean/mujoco/mjc_half_cheetah.c create mode 100644 ocean/mujoco/mjc_half_cheetah.cu create mode 100644 ocean/mujoco/mjc_half_cheetah.h create mode 100644 ocean/mujoco/mjc_hopper.c create mode 100644 ocean/mujoco/mjc_hopper.cu create mode 100644 ocean/mujoco/mjc_hopper.h create mode 100644 ocean/mujoco/mjc_humanoid.c create mode 100644 ocean/mujoco/mjc_humanoid.cu create mode 100644 ocean/mujoco/mjc_humanoid.h create mode 100644 ocean/mujoco/mjc_walker2d.c create mode 100644 ocean/mujoco/mjc_walker2d.cu create mode 100644 ocean/mujoco/mjc_walker2d.h create mode 100644 ocean/mujoco/mjcf2bin.py create mode 100644 ocean/mujoco/physics.h create mode 100644 ocean/mujoco/render.h create mode 100644 resources/mujoco/ant.xml create mode 100644 resources/mujoco/half_cheetah.xml create mode 100644 resources/mujoco/hopper.xml create mode 100644 resources/mujoco/humanoid.xml create mode 100644 resources/mujoco/walker2d.xml diff --git a/.gitignore b/.gitignore index 8195614174..4e469490d8 100644 --- a/.gitignore +++ b/.gitignore @@ -178,6 +178,7 @@ resources/drive/data/* resources/drive/binaries/* resources/boxoban/levels/ resources/boxoban/boxoban_maps_*.bin +resources/mujoco/*.bin # Policy weights live in the website repo (docs/assets/models/) resources/**/*_weights.bin diff --git a/build.sh b/build.sh index 115d8f4b73..80f96866d4 100755 --- a/build.sh +++ b/build.sh @@ -13,6 +13,8 @@ set -e # ./build.sh breakout --web # Emscripten web build # # copy build/web/ENV/* to ../docker/puffer.ai/docs/assets/ENV/ # ./build.sh breakout --profile # Kernel profiling binary +# ./build.sh mjc_half_cheetah # MuJoCo-family envs live in ocean/mujoco/ENV.h +# ./build.sh mjc_half_cheetah --cu # ... and ocean/mujoco/ENV.cu (one GPU thread per env) # ./build.sh all # Build all envs native and native float32 # # Env is compiled in. Run: ./puffer train|eval|match|sweep [--section.key=value ...] @@ -182,6 +184,8 @@ elif [ "$ENV" = "nethack" ]; then -Xlinker -rpath -Xlinker "$NETHACK_LIB_DIR" -ldl) elif [ -d "ocean/$ENV" ]; then SRC_DIR="ocean/$ENV" +elif [ -f "ocean/mujoco/$ENV.h" ]; then + SRC_DIR="ocean/mujoco" else echo "Error: environment '$ENV' not found" && exit 1 fi diff --git a/config/mjc_ant.ini b/config/mjc_ant.ini new file mode 100644 index 0000000000..87dd196491 --- /dev/null +++ b/config/mjc_ant.ini @@ -0,0 +1,48 @@ +[base] +env_name = mjc_ant + +[vec] +total_agents = 2048 +# GPU env builds (--cu) require 1 buffer; the CPU vec is ~1.5x faster with 4 +num_buffers = 1 +num_threads = 16 + +[env] +model = resources/mujoco/ant.bin +# Gymnasium Ant-v5 defaults +max_steps = 1000 +reset_noise_scale = 0.1 +forward_reward_weight = 1.0 +ctrl_cost_weight = 0.5 +contact_cost_weight = 0.0005 +healthy_reward = 1.0 + +[policy] +hidden_size = 256 +num_layers = 2 +expansion_factor = 1 + +[train] +gpus = 1 +seed = 42 +total_timesteps = 100000000 +learning_rate = 0.005 +anneal_lr = 1 +min_lr_ratio = 0 +gamma = 0.99 +gae_lambda = 0.95 +replay_ratio = 2 +clip_coef = 0.2 +vf_coef = 0.5 +vf_clip_coef = 1.0 +max_grad_norm = 0.5 +ent_coef = 0.0 +momentum = 0.9 +minibatch_size = 8192 +horizon = 64 +vtrace_rho_clip = 1.0 +vtrace_c_clip = 1.0 + +[sweep] +metric = score +goal = maximize diff --git a/config/mjc_half_cheetah.ini b/config/mjc_half_cheetah.ini new file mode 100644 index 0000000000..e7978c813e --- /dev/null +++ b/config/mjc_half_cheetah.ini @@ -0,0 +1,46 @@ +[base] +env_name = mjc_half_cheetah + +[vec] +total_agents = 2048 +# GPU env builds (--cu) require 1 buffer; the CPU vec is ~1.5x faster with 4 +num_buffers = 1 +num_threads = 16 + +[env] +model = resources/mujoco/half_cheetah.bin +# Gymnasium HalfCheetah-v5 defaults. Episode ends (terminal) at max_steps. +max_steps = 1000 +reset_noise_scale = 0.1 +forward_reward_weight = 1.0 +ctrl_cost_weight = 0.1 + +[policy] +hidden_size = 256 +num_layers = 2 +expansion_factor = 1 + +[train] +gpus = 1 +seed = 42 +total_timesteps = 100000000 +learning_rate = 0.005 +anneal_lr = 1 +min_lr_ratio = 0 +gamma = 0.99 +gae_lambda = 0.95 +replay_ratio = 2 +clip_coef = 0.2 +vf_coef = 0.5 +vf_clip_coef = 1.0 +max_grad_norm = 0.5 +ent_coef = 0.0 +momentum = 0.9 +minibatch_size = 8192 +horizon = 64 +vtrace_rho_clip = 1.0 +vtrace_c_clip = 1.0 + +[sweep] +metric = score +goal = maximize diff --git a/config/mjc_hopper.ini b/config/mjc_hopper.ini new file mode 100644 index 0000000000..e256ed8120 --- /dev/null +++ b/config/mjc_hopper.ini @@ -0,0 +1,47 @@ +[base] +env_name = mjc_hopper + +[vec] +total_agents = 2048 +# GPU env builds (--cu) require 1 buffer; the CPU vec is ~1.5x faster with 4 +num_buffers = 1 +num_threads = 16 + +[env] +model = resources/mujoco/hopper.bin +# Gymnasium Hopper-v5 defaults +max_steps = 1000 +reset_noise_scale = 0.005 +forward_reward_weight = 1.0 +ctrl_cost_weight = 0.001 +healthy_reward = 1.0 + +[policy] +hidden_size = 256 +num_layers = 2 +expansion_factor = 1 + +[train] +gpus = 1 +seed = 42 +total_timesteps = 100000000 +learning_rate = 0.005 +anneal_lr = 1 +min_lr_ratio = 0 +gamma = 0.99 +gae_lambda = 0.95 +replay_ratio = 2 +clip_coef = 0.2 +vf_coef = 0.5 +vf_clip_coef = 1.0 +max_grad_norm = 0.5 +ent_coef = 0.0 +momentum = 0.9 +minibatch_size = 8192 +horizon = 64 +vtrace_rho_clip = 1.0 +vtrace_c_clip = 1.0 + +[sweep] +metric = score +goal = maximize diff --git a/config/mjc_humanoid.ini b/config/mjc_humanoid.ini new file mode 100644 index 0000000000..8d78c85ea9 --- /dev/null +++ b/config/mjc_humanoid.ini @@ -0,0 +1,48 @@ +[base] +env_name = mjc_humanoid + +[vec] +total_agents = 2048 +# GPU env builds (--cu) require 1 buffer; the CPU vec is ~1.5x faster with 4 +num_buffers = 1 +num_threads = 16 + +[env] +model = resources/mujoco/humanoid.bin +# Gymnasium Humanoid-v5 defaults (contact cost clamped at 10 in the env) +max_steps = 1000 +reset_noise_scale = 0.01 +forward_reward_weight = 1.25 +ctrl_cost_weight = 0.1 +contact_cost_weight = 0.0000005 +healthy_reward = 5.0 + +[policy] +hidden_size = 256 +num_layers = 2 +expansion_factor = 1 + +[train] +gpus = 1 +seed = 42 +total_timesteps = 100000000 +learning_rate = 0.005 +anneal_lr = 1 +min_lr_ratio = 0 +gamma = 0.99 +gae_lambda = 0.95 +replay_ratio = 2 +clip_coef = 0.2 +vf_coef = 0.5 +vf_clip_coef = 1.0 +max_grad_norm = 0.5 +ent_coef = 0.0 +momentum = 0.9 +minibatch_size = 8192 +horizon = 64 +vtrace_rho_clip = 1.0 +vtrace_c_clip = 1.0 + +[sweep] +metric = score +goal = maximize diff --git a/config/mjc_walker2d.ini b/config/mjc_walker2d.ini new file mode 100644 index 0000000000..fad1459c43 --- /dev/null +++ b/config/mjc_walker2d.ini @@ -0,0 +1,47 @@ +[base] +env_name = mjc_walker2d + +[vec] +total_agents = 2048 +# GPU env builds (--cu) require 1 buffer; the CPU vec is ~1.5x faster with 4 +num_buffers = 1 +num_threads = 16 + +[env] +model = resources/mujoco/walker2d.bin +# Gymnasium Walker2d-v5 defaults +max_steps = 1000 +reset_noise_scale = 0.005 +forward_reward_weight = 1.0 +ctrl_cost_weight = 0.001 +healthy_reward = 1.0 + +[policy] +hidden_size = 256 +num_layers = 2 +expansion_factor = 1 + +[train] +gpus = 1 +seed = 42 +total_timesteps = 100000000 +learning_rate = 0.005 +anneal_lr = 1 +min_lr_ratio = 0 +gamma = 0.99 +gae_lambda = 0.95 +replay_ratio = 2 +clip_coef = 0.2 +vf_coef = 0.5 +vf_clip_coef = 1.0 +max_grad_norm = 0.5 +ent_coef = 0.0 +momentum = 0.9 +minibatch_size = 8192 +horizon = 64 +vtrace_rho_clip = 1.0 +vtrace_c_clip = 1.0 + +[sweep] +metric = score +goal = maximize diff --git a/ocean/mujoco/backend.h b/ocean/mujoco/backend.h new file mode 100644 index 0000000000..2d3aab72ae --- /dev/null +++ b/ocean/mujoco/backend.h @@ -0,0 +1,129 @@ +// Vector backends for the ocean/mujoco envs, included at the end of each +// mjc_.h after Env, mjc_reset, mjc_step, mjc_render and mjc_init. All +// memory is bound at create time: the solver scratch (MJ_SCRATCH floats per +// env, see mj_makeData) and, on the GPU, the whole batch. CPU: the puf_* per-env +// API calls straight through. GPU (mjc_.cu defines PUF_BACKEND PUF_GPU): +// one thread per env steps a thread-local Env (local memory is lane interleaved, +// so every access coalesces) and copies only the persistent prefix, everything +// before MjData.xquat, in and out of the device batch; the scratch is laid out +// lane interleaved too (element stride 32). The trainer's device obs/action/ +// reward/terminal buffers are bound into Env.agents at create time. + +#if PUF_BACKEND == PUF_GPU +#define MJC_BLOCK 128 +#define MJC_STATE (offsetof(Env, d) + offsetof(MjData, xquat)) + +struct { + Env* envs; + int n; + cudaStream_t stream; +} mjc_gpu; + +__global__ void mjc_reset_kernel(Env* envs, int n) { + int i = blockIdx.x*blockDim.x + threadIdx.x; + if (i < n) { + Env env; + memcpy(&env, &envs[i], MJC_STATE); + mjc_reset(&env); + memcpy(&envs[i], &env, MJC_STATE); + } +} + +__global__ void mjc_step_kernel(Env* envs, int n) { + int i = blockIdx.x*blockDim.x + threadIdx.x; + if (i < n) { + Env env; + memcpy(&env, &envs[i], MJC_STATE); + mjc_step(&env); + memcpy(&envs[i], &env, MJC_STATE); + } +} + +Env* puf_vec_create(int n, Dict* kwargs, obs_t* observations, float* actions, float* rewards, + float* terminals) { + Env* host = (Env*)calloc(n, sizeof(Env)); + MjModel* m; + assert(cudaMalloc((void**)&m, sizeof(MjModel)) == cudaSuccess); + float* scratch; + size_t groups = (n + 31) / 32; + assert(cudaMalloc((void**)&scratch, groups*32*MJ_SCRATCH*sizeof(float)) == cudaSuccess + && "GPU env solver scratch does not fit in device memory"); + for (int i = 0; i < n; i++) { + Env* env = &host[i]; + mjc_init(env, kwargs); + mj_makeData(&env->d, scratch + (size_t)(i / 32)*32*MJ_SCRATCH + i % 32, 32); + env->m = m; + env->rng = i + 1; + env->agents[0].observations = observations + (long)i*OBS_SIZE; + env->agents[0].actions = actions + (long)i*NUM_ATNS; + env->agents[0].rewards = rewards + i; + env->agents[0].terminals = terminals + i; + } + cudaMemcpy(m, &mj_model, sizeof(MjModel), cudaMemcpyHostToDevice); + assert(cudaMalloc((void**)&mjc_gpu.envs, (size_t)n*sizeof(Env)) == cudaSuccess + && "GPU env batch does not fit in device memory"); + cudaMemcpy(mjc_gpu.envs, host, (size_t)n*sizeof(Env), cudaMemcpyHostToDevice); + free(host); + mjc_gpu.n = n; + return mjc_gpu.envs; +} + +void puf_bind_stream(cudaStream_t stream) { + mjc_gpu.stream = stream; +} + +void puf_init(Env* env, Dict* kwargs) { +} + +void puf_reset(Env* envs) { + mjc_reset_kernel<<<(mjc_gpu.n + MJC_BLOCK - 1) / MJC_BLOCK, MJC_BLOCK>>>(mjc_gpu.envs, + mjc_gpu.n); + assert(cudaGetLastError() == cudaSuccess); +} + +void puf_step(Env* envs) { + mjc_step_kernel<<<(mjc_gpu.n + MJC_BLOCK - 1) / MJC_BLOCK, MJC_BLOCK, 0, mjc_gpu.stream>>>( + mjc_gpu.envs, mjc_gpu.n); + assert(cudaGetLastError() == cudaSuccess); +} + +// Copy env 0's state back and draw it with the host model (mj_render recomputes +// kinematics and contacts from qpos) +void puf_render(Env* envs) { + Env env; + cudaStreamSynchronize(mjc_gpu.stream); + cudaMemcpy(&env, mjc_gpu.envs, sizeof(Env), cudaMemcpyDeviceToHost); + env.m = &mj_model; + mjc_render(&env); +} + +void puf_close(Env* envs) { + cudaFree(mjc_gpu.envs); + if (IsWindowReady()) { + CloseWindow(); + } +} +#else +void puf_init(Env* env, Dict* kwargs) { + mjc_init(env, kwargs); + mj_makeData(&env->d, (float*)calloc(MJ_SCRATCH, sizeof(float)), 1); +} + +void puf_reset(Env* env) { + mjc_reset(env); +} + +void puf_step(Env* env) { + mjc_step(env); +} + +void puf_render(Env* env) { + mjc_render(env); +} + +void puf_close(Env* env) { + if (IsWindowReady()) { + CloseWindow(); + } +} +#endif diff --git a/ocean/mujoco/mjc_ant.c b/ocean/mujoco/mjc_ant.c new file mode 100644 index 0000000000..55b10a3bb3 --- /dev/null +++ b/ocean/mujoco/mjc_ant.c @@ -0,0 +1,3 @@ +#define PUFFERCPU_EVAL_MAIN +#define ENV_HEADER "../ocean/mujoco/mjc_ant.h" +#include "puffercpu.h" diff --git a/ocean/mujoco/mjc_ant.cu b/ocean/mujoco/mjc_ant.cu new file mode 100644 index 0000000000..045781e34a --- /dev/null +++ b/ocean/mujoco/mjc_ant.cu @@ -0,0 +1,4 @@ +// GPU build of mjc_ant: one thread per env (see backend.h) +#define PUF_BACKEND PUF_GPU +typedef float obs_t; +#include "mjc_ant.h" diff --git a/ocean/mujoco/mjc_ant.h b/ocean/mujoco/mjc_ant.h new file mode 100644 index 0000000000..a0bc7e5261 --- /dev/null +++ b/ocean/mujoco/mjc_ant.h @@ -0,0 +1,176 @@ +// Ant (gymnasium Ant-v5) on the MuJoCo-style physics core: obs = qpos[2:] + +// qvel + clip(cfrc_ext[1:], -1, 1), reward = healthy + forward velocity - ctrl +// cost - contact cost, terminates when the torso leaves [0.2, 1.0] m. Model: +// resources/mujoco/ant.xml compiled by mjcf2bin.py. mjc_reset/mjc_step are +// host+device (see backend.h). + +#include +#include +#include +#include +#include "raylib.h" +typedef float obs_t; +#include "pufferenv.h" +// physics.h capacities sized to this model (checked by mj_loadModel) +#define MJ_MAX_NQ 15 +#define MJ_MAX_NV 14 +#define MJ_MAX_NBODY 14 +#define MJ_MAX_NJNT 9 +#define MJ_MAX_NGEOM 14 +#define MJ_MAX_NU 8 +#define MJ_MAXCON 25 +#define MJ_MAXEFC 108 +#include "physics.h" +#include "render.h" + +#define ANT_FRAME_SKIP 5 +#define OBS_SIZE 105 +#define NUM_ATNS 8 +#define ACT_SIZES {1, 1, 1, 1, 1, 1, 1, 1} +#define PUF_STEPS_PER_SEC 20 + +MjModel mj_model; + +struct Log { + float perf; + float score; + float episode_return; + float episode_length; + float x_velocity; + float distance; + float n; +}; + +struct Env { + Log log; + Agent agents[1]; + int tag; + int boundary_reached; + int num_agents; + unsigned int rng; + const MjModel* m; + int tick; + float x_start; + float episode_return; + int max_steps; + float reset_noise_scale; + float forward_reward_weight; + float ctrl_cost_weight; + float contact_cost_weight; + float healthy_reward; + MjData d; +}; +typedef Env Ant; + +// Returns the contact cost (sum of squared clipped external forces) +MJ_HD float compute_observations(Ant* env) { + const MjModel* m = env->m; + float* obs = env->agents[0].observations; + memcpy(obs, env->d.qpos + 2, (m->nq - 2)*sizeof(float)); + memcpy(obs + m->nq - 2, env->d.qvel, m->nv*sizeof(float)); + float* cfrc = obs + m->nq - 2 + m->nv; + float cost = 0.0f; + for (int b = 1; b < m->nbody; b++) { + for (int k = 0; k < 6; k++) { + float f = fminf(fmaxf(env->d.cfrc_ext[b][k], -1.0f), 1.0f); + cfrc[6*(b - 1) + k] = f; + cost += f*f; + } + } + return env->contact_cost_weight*cost; +} + +MJ_HD void mjc_reset(Ant* env) { + const MjModel* m = env->m; + mj_resetData(m, &env->d); + float s = env->reset_noise_scale; + for (int i = 0; i < m->nq; i++) { + env->d.qpos[i] += s*(2.0f*mju_rand(&env->rng) - 1.0f); + } + for (int i = 0; i < m->nv; i++) { + env->d.qvel[i] = s*mju_randn(&env->rng); + } + mj_kinematics(m, &env->d); + env->tick = 0; + env->x_start = env->d.qpos[0]; + env->episode_return = 0.0f; + compute_observations(env); +} + +MJ_HD void mjc_step(Ant* env) { + const MjModel* m = env->m; + float* actions = env->agents[0].actions; + float cost = 0.0f; + for (int i = 0; i < m->nu; i++) { + env->d.ctrl[i] = fminf(fmaxf(actions[i], -1.0f), 1.0f); + cost += env->d.ctrl[i]*env->d.ctrl[i]; + } + // Gym measures the torso displacement with body xpos, which lags qpos by + // one substep (kinematics of the last forward pass) + float x0 = env->d.xpos[1][0]; + for (int k = 0; k < ANT_FRAME_SKIP; k++) { + mj_step(m, &env->d); + } + mj_rnePostConstraint(m, &env->d); + float* q = env->d.qpos; + float dt = ANT_FRAME_SKIP*m->opt_timestep; + float x_velocity = (env->d.xpos[1][0] - x0) / dt; + int healthy = isfinite(q[2]) && q[2] >= 0.2f && q[2] <= 1.0f; + float contact_cost = compute_observations(env); + float reward = env->forward_reward_weight*x_velocity - env->ctrl_cost_weight*cost + - contact_cost + (healthy ? env->healthy_reward : 0.0f); + env->tick++; + env->episode_return += reward; + env->agents[0].rewards[0] = reward; + env->agents[0].terminals[0] = 0.0f; + if (healthy && env->tick < env->max_steps) { + return; + } + float distance = q[0] - env->x_start; + float xvel = distance / (env->tick*dt); + env->agents[0].terminals[0] = 1.0f; + env->log.perf += fminf(fmaxf(xvel / 5.0f, 0.0f), 1.0f); + env->log.score += env->episode_return; + env->log.episode_return += env->episode_return; + env->log.episode_length += env->tick; + env->log.x_velocity += xvel; + env->log.distance += distance; + env->log.n += 1.0f; + mjc_reset(env); +} + +void mjc_render(Ant* env) { + float target[3] = {env->d.xpos[1][0], env->d.xpos[1][1], 0.3f}; + mj_render(env->m, &env->d, "PufferLib Ant", target, 4.0f, 2.0f, + TextFormat("step %d x %.2f m vel %.2f m/s return %.1f", env->tick, + env->d.qpos[0], env->d.qvel[0], env->episode_return)); +} + +void mjc_init(Ant* env, Dict* kwargs) { + if (mj_model.nbody == 0) { + mj_loadModel(&mj_model, dict_get_str(kwargs, "model")); + } + assert(mj_model.nq - 2 + mj_model.nv + 6*(mj_model.nbody - 1) == OBS_SIZE); + assert(mj_model.nu == NUM_ATNS); + env->m = &mj_model; + env->num_agents = 1; + env->agents[0].policy = 0; + env->agents[0].action_mask = NULL; + env->max_steps = dict_get(kwargs, "max_steps"); + env->reset_noise_scale = dict_get(kwargs, "reset_noise_scale"); + env->forward_reward_weight = dict_get(kwargs, "forward_reward_weight"); + env->ctrl_cost_weight = dict_get(kwargs, "ctrl_cost_weight"); + env->contact_cost_weight = dict_get(kwargs, "contact_cost_weight"); + env->healthy_reward = dict_get(kwargs, "healthy_reward"); +} + +void puf_log(Log* log, Dict* out) { + dict_set(out, "perf", log->perf); + dict_set(out, "score", log->score); + dict_set(out, "episode_return", log->episode_return); + dict_set(out, "episode_length", log->episode_length); + dict_set(out, "x_velocity", log->x_velocity); + dict_set(out, "distance", log->distance); +} + +#include "backend.h" diff --git a/ocean/mujoco/mjc_half_cheetah.c b/ocean/mujoco/mjc_half_cheetah.c new file mode 100644 index 0000000000..d02fd4d5f1 --- /dev/null +++ b/ocean/mujoco/mjc_half_cheetah.c @@ -0,0 +1,3 @@ +#define PUFFERCPU_EVAL_MAIN +#define ENV_HEADER "../ocean/mujoco/mjc_half_cheetah.h" +#include "puffercpu.h" diff --git a/ocean/mujoco/mjc_half_cheetah.cu b/ocean/mujoco/mjc_half_cheetah.cu new file mode 100644 index 0000000000..e7329f49a9 --- /dev/null +++ b/ocean/mujoco/mjc_half_cheetah.cu @@ -0,0 +1,4 @@ +// GPU build of mjc_half_cheetah: one thread per env (see backend.h) +#define PUF_BACKEND PUF_GPU +typedef float obs_t; +#include "mjc_half_cheetah.h" diff --git a/ocean/mujoco/mjc_half_cheetah.h b/ocean/mujoco/mjc_half_cheetah.h new file mode 100644 index 0000000000..158f912d74 --- /dev/null +++ b/ocean/mujoco/mjc_half_cheetah.h @@ -0,0 +1,158 @@ +// HalfCheetah (gymnasium HalfCheetah-v5) on the MuJoCo-style physics core: +// obs = qpos[1:] + qvel, reward = forward velocity - ctrl cost, 1000 step +// episodes. Model: resources/mujoco/half_cheetah.xml compiled by mjcf2bin.py. +// mjc_reset/mjc_step are host+device; backend.h wraps them for the CPU vec +// or launches them one thread per env on the GPU (--cu). + +#include +#include +#include +#include +#include "raylib.h" +typedef float obs_t; +#include "pufferenv.h" +// physics.h capacities sized to this model (checked by mj_loadModel) +#define MJ_MAX_NQ 9 +#define MJ_MAX_NV 9 +#define MJ_MAX_NBODY 8 +#define MJ_MAX_NJNT 9 +#define MJ_MAX_NGEOM 9 +#define MJ_MAX_NU 6 +#define MJ_MAXCON 16 +#define MJ_MAXEFC 70 +#include "physics.h" +#include "render.h" + +#define HC_FRAME_SKIP 5 +#define OBS_SIZE 17 +#define NUM_ATNS 6 +#define ACT_SIZES {1, 1, 1, 1, 1, 1} +#define PUF_STEPS_PER_SEC 20 + +MjModel mj_model; + +struct Log { + float perf; + float score; + float episode_return; + float episode_length; + float x_velocity; + float distance; + float ctrl_cost; + float n; +}; + +struct Env { + Log log; + Agent agents[1]; + int tag; + int boundary_reached; + int num_agents; + unsigned int rng; + const MjModel* m; + int tick; + float x_start; + float episode_return; + float episode_ctrl_cost; + int max_steps; + float reset_noise_scale; + float forward_reward_weight; + float ctrl_cost_weight; + MjData d; +}; +typedef Env HalfCheetah; + +MJ_HD void compute_observations(HalfCheetah* env) { + float* obs = env->agents[0].observations; + memcpy(obs, env->d.qpos + 1, (env->m->nq - 1)*sizeof(float)); + memcpy(obs + env->m->nq - 1, env->d.qvel, env->m->nv*sizeof(float)); +} + +MJ_HD void mjc_reset(HalfCheetah* env) { + const MjModel* m = env->m; + mj_resetData(m, &env->d); + float s = env->reset_noise_scale; + for (int i = 0; i < m->nq; i++) { + env->d.qpos[i] += s*(2.0f*mju_rand(&env->rng) - 1.0f); + } + for (int i = 0; i < m->nv; i++) { + env->d.qvel[i] = s*mju_randn(&env->rng); + } + env->tick = 0; + env->x_start = env->d.qpos[0]; + env->episode_return = 0.0f; + env->episode_ctrl_cost = 0.0f; + compute_observations(env); +} + +MJ_HD void mjc_step(HalfCheetah* env) { + const MjModel* m = env->m; + float* actions = env->agents[0].actions; + float cost = 0.0f; + for (int i = 0; i < m->nu; i++) { + env->d.ctrl[i] = fminf(fmaxf(actions[i], -1.0f), 1.0f); + cost += env->d.ctrl[i]*env->d.ctrl[i]; + } + float x0 = env->d.qpos[0]; + for (int k = 0; k < HC_FRAME_SKIP; k++) { + mj_step(m, &env->d); + } + float dt = HC_FRAME_SKIP*m->opt_timestep; + float x_velocity = (env->d.qpos[0] - x0) / dt; + float reward = env->forward_reward_weight*x_velocity - env->ctrl_cost_weight*cost; + env->tick++; + env->episode_return += reward; + env->episode_ctrl_cost += env->ctrl_cost_weight*cost; + env->agents[0].rewards[0] = reward; + env->agents[0].terminals[0] = 0.0f; + if (env->tick < env->max_steps) { + compute_observations(env); + return; + } + float distance = env->d.qpos[0] - env->x_start; + float xvel = distance / (env->tick*dt); + env->agents[0].terminals[0] = 1.0f; + env->log.perf += fminf(fmaxf(xvel / 10.0f, 0.0f), 1.0f); + env->log.score += env->episode_return; + env->log.episode_return += env->episode_return; + env->log.episode_length += env->tick; + env->log.x_velocity += xvel; + env->log.distance += distance; + env->log.ctrl_cost += env->episode_ctrl_cost; + env->log.n += 1.0f; + mjc_reset(env); +} + +void mjc_render(HalfCheetah* env) { + float target[3] = {env->d.xpos[1][0], 0.0f, 0.5f}; + mj_render(env->m, &env->d, "PufferLib HalfCheetah", target, 4.0f, 1.0f, + TextFormat("step %d x %.2f m vel %.2f m/s return %.1f", env->tick, + env->d.qpos[0], env->d.qvel[0], env->episode_return)); +} + +void mjc_init(HalfCheetah* env, Dict* kwargs) { + if (mj_model.nbody == 0) { + mj_loadModel(&mj_model, dict_get_str(kwargs, "model")); + } + assert(mj_model.nq - 1 + mj_model.nv == OBS_SIZE && mj_model.nu == NUM_ATNS); + env->m = &mj_model; + env->num_agents = 1; + env->agents[0].policy = 0; + env->agents[0].action_mask = NULL; + env->max_steps = dict_get(kwargs, "max_steps"); + env->reset_noise_scale = dict_get(kwargs, "reset_noise_scale"); + env->forward_reward_weight = dict_get(kwargs, "forward_reward_weight"); + env->ctrl_cost_weight = dict_get(kwargs, "ctrl_cost_weight"); +} + +void puf_log(Log* log, Dict* out) { + dict_set(out, "perf", log->perf); + dict_set(out, "score", log->score); + dict_set(out, "episode_return", log->episode_return); + dict_set(out, "episode_length", log->episode_length); + dict_set(out, "x_velocity", log->x_velocity); + dict_set(out, "distance", log->distance); + dict_set(out, "ctrl_cost", log->ctrl_cost); +} + +#include "backend.h" diff --git a/ocean/mujoco/mjc_hopper.c b/ocean/mujoco/mjc_hopper.c new file mode 100644 index 0000000000..1aff4db20b --- /dev/null +++ b/ocean/mujoco/mjc_hopper.c @@ -0,0 +1,3 @@ +#define PUFFERCPU_EVAL_MAIN +#define ENV_HEADER "../ocean/mujoco/mjc_hopper.h" +#include "puffercpu.h" diff --git a/ocean/mujoco/mjc_hopper.cu b/ocean/mujoco/mjc_hopper.cu new file mode 100644 index 0000000000..ada27027f5 --- /dev/null +++ b/ocean/mujoco/mjc_hopper.cu @@ -0,0 +1,4 @@ +// GPU build of mjc_hopper: one thread per env (see backend.h) +#define PUF_BACKEND PUF_GPU +typedef float obs_t; +#include "mjc_hopper.h" diff --git a/ocean/mujoco/mjc_hopper.h b/ocean/mujoco/mjc_hopper.h new file mode 100644 index 0000000000..3eff17b746 --- /dev/null +++ b/ocean/mujoco/mjc_hopper.h @@ -0,0 +1,165 @@ +// Hopper (gymnasium Hopper-v5) on the MuJoCo-style physics core: obs = +// qpos[1:] + clip(qvel, -10, 10), reward = healthy + forward velocity - ctrl +// cost, terminates when unhealthy. Model: resources/mujoco/hopper.xml compiled +// by mjcf2bin.py. mjc_reset/mjc_step are host+device (see backend.h). + +#include +#include +#include +#include +#include "raylib.h" +typedef float obs_t; +#include "pufferenv.h" +// physics.h capacities sized to this model +#define MJ_MAX_NQ 6 +#define MJ_MAX_NV 6 +#define MJ_MAX_NBODY 5 +#define MJ_MAX_NJNT 6 +#define MJ_MAX_NGEOM 5 +#define MJ_MAX_NU 3 +#define MJ_MAXCON 8 +#define MJ_MAXEFC 35 +#include "physics.h" +#include "render.h" + +#define HP_FRAME_SKIP 4 +#define OBS_SIZE 11 +#define NUM_ATNS 3 +#define ACT_SIZES {1, 1, 1} +#define PUF_STEPS_PER_SEC 125 + +MjModel mj_model; + +struct Log { + float perf; + float score; + float episode_return; + float episode_length; + float x_velocity; + float distance; + float n; +}; + +struct Env { + Log log; + Agent agents[1]; + int tag; + int boundary_reached; + int num_agents; + unsigned int rng; + const MjModel* m; + int tick; + float x_start; + float episode_return; + int max_steps; + float reset_noise_scale; + float forward_reward_weight; + float ctrl_cost_weight; + float healthy_reward; + MjData d; +}; +typedef Env Hopper; + +MJ_HD void compute_observations(Hopper* env) { + const MjModel* m = env->m; + float* obs = env->agents[0].observations; + memcpy(obs, env->d.qpos + 1, (m->nq - 1)*sizeof(float)); + for (int i = 0; i < m->nv; i++) { + obs[m->nq - 1 + i] = fminf(fmaxf(env->d.qvel[i], -10.0f), 10.0f); + } +} + +MJ_HD void mjc_reset(Hopper* env) { + const MjModel* m = env->m; + mj_resetData(m, &env->d); + float s = env->reset_noise_scale; + for (int i = 0; i < m->nq; i++) { + env->d.qpos[i] += s*(2.0f*mju_rand(&env->rng) - 1.0f); + } + for (int i = 0; i < m->nv; i++) { + env->d.qvel[i] = s*(2.0f*mju_rand(&env->rng) - 1.0f); + } + env->tick = 0; + env->x_start = env->d.qpos[0]; + env->episode_return = 0.0f; + compute_observations(env); +} + +MJ_HD void mjc_step(Hopper* env) { + const MjModel* m = env->m; + float* actions = env->agents[0].actions; + float cost = 0.0f; + for (int i = 0; i < m->nu; i++) { + env->d.ctrl[i] = fminf(fmaxf(actions[i], -1.0f), 1.0f); + cost += env->d.ctrl[i]*env->d.ctrl[i]; + } + float x0 = env->d.qpos[0]; + for (int k = 0; k < HP_FRAME_SKIP; k++) { + mj_step(m, &env->d); + } + float* q = env->d.qpos; + float dt = HP_FRAME_SKIP*m->opt_timestep; + float x_velocity = (q[0] - x0) / dt; + int healthy = q[1] > 0.7f && q[2] > -0.2f && q[2] < 0.2f; + for (int i = 2; i < m->nq; i++) { + healthy = healthy && q[i] > -100.0f && q[i] < 100.0f; + } + for (int i = 0; i < m->nv; i++) { + healthy = healthy && env->d.qvel[i] > -100.0f && env->d.qvel[i] < 100.0f; + } + float reward = env->forward_reward_weight*x_velocity - env->ctrl_cost_weight*cost + + (healthy ? env->healthy_reward : 0.0f); + env->tick++; + env->episode_return += reward; + env->agents[0].rewards[0] = reward; + env->agents[0].terminals[0] = 0.0f; + if (healthy && env->tick < env->max_steps) { + compute_observations(env); + return; + } + float distance = q[0] - env->x_start; + float xvel = distance / (env->tick*dt); + env->agents[0].terminals[0] = 1.0f; + env->log.perf += fminf(fmaxf(xvel / 3.0f, 0.0f), 1.0f); + env->log.score += env->episode_return; + env->log.episode_return += env->episode_return; + env->log.episode_length += env->tick; + env->log.x_velocity += xvel; + env->log.distance += distance; + env->log.n += 1.0f; + mjc_reset(env); +} + +void mjc_render(Hopper* env) { + float target[3] = {env->d.xpos[1][0], 0.0f, 0.8f}; + mj_render(env->m, &env->d, "PufferLib Hopper", target, 4.0f, 1.0f, + TextFormat("step %d x %.2f m vel %.2f m/s return %.1f", env->tick, + env->d.qpos[0], env->d.qvel[0], env->episode_return)); +} + +void mjc_init(Hopper* env, Dict* kwargs) { + if (mj_model.nbody == 0) { + mj_loadModel(&mj_model, dict_get_str(kwargs, "model")); + } + assert(mj_model.nq - 1 + mj_model.nv == OBS_SIZE && mj_model.nu == NUM_ATNS); + env->m = &mj_model; + env->num_agents = 1; + env->agents[0].policy = 0; + env->agents[0].action_mask = NULL; + env->max_steps = dict_get(kwargs, "max_steps"); + env->reset_noise_scale = dict_get(kwargs, "reset_noise_scale"); + env->forward_reward_weight = dict_get(kwargs, "forward_reward_weight"); + env->ctrl_cost_weight = dict_get(kwargs, "ctrl_cost_weight"); + env->healthy_reward = dict_get(kwargs, "healthy_reward"); +} + +void puf_log(Log* log, Dict* out) { + dict_set(out, "perf", log->perf); + dict_set(out, "score", log->score); + dict_set(out, "episode_return", log->episode_return); + dict_set(out, "episode_length", log->episode_length); + dict_set(out, "x_velocity", log->x_velocity); + dict_set(out, "distance", log->distance); +} + +#include "backend.h" diff --git a/ocean/mujoco/mjc_humanoid.c b/ocean/mujoco/mjc_humanoid.c new file mode 100644 index 0000000000..3873e82fef --- /dev/null +++ b/ocean/mujoco/mjc_humanoid.c @@ -0,0 +1,3 @@ +#define PUFFERCPU_EVAL_MAIN +#define ENV_HEADER "../ocean/mujoco/mjc_humanoid.h" +#include "puffercpu.h" diff --git a/ocean/mujoco/mjc_humanoid.cu b/ocean/mujoco/mjc_humanoid.cu new file mode 100644 index 0000000000..ad04a9d4ce --- /dev/null +++ b/ocean/mujoco/mjc_humanoid.cu @@ -0,0 +1,4 @@ +// GPU build of mjc_humanoid: one thread per env (see backend.h) +#define PUF_BACKEND PUF_GPU +typedef float obs_t; +#include "mjc_humanoid.h" diff --git a/ocean/mujoco/mjc_humanoid.h b/ocean/mujoco/mjc_humanoid.h new file mode 100644 index 0000000000..0d3e9eaff1 --- /dev/null +++ b/ocean/mujoco/mjc_humanoid.h @@ -0,0 +1,194 @@ +// Humanoid (gymnasium Humanoid-v5) on the MuJoCo-style physics core: obs = +// qpos[2:] + qvel + cinert[1:] + cvel[1:] + qfrc_actuator[6:] + cfrc_ext[1:], +// reward = healthy + forward COM velocity - ctrl cost - contact cost (clamped +// at 10), terminates when the torso leaves z in (1, 2). Model: +// resources/mujoco/humanoid.xml compiled by mjcf2bin.py. mjc_reset/mjc_step +// are host+device (see backend.h). + +#include +#include +#include +#include +#include "raylib.h" +typedef float obs_t; +#include "pufferenv.h" +// physics.h capacities sized to this model (checked by mj_loadModel) +#define MJ_MAX_NQ 24 +#define MJ_MAX_NV 23 +#define MJ_MAX_NBODY 14 +#define MJ_MAX_NJNT 18 +#define MJ_MAX_NGEOM 18 +#define MJ_MAX_NU 17 +#define MJ_MAXCON 32 +#define MJ_MAXEFC 145 +#include "physics.h" +#include "render.h" + +#define HM_FRAME_SKIP 5 +#define HM_CONTACT_COST_MAX 10.0f +#define OBS_SIZE 348 +#define NUM_ATNS 17 +#define ACT_SIZES {1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1} +#define PUF_STEPS_PER_SEC 67 + +MjModel mj_model; + +struct Log { + float perf; + float score; + float episode_return; + float episode_length; + float x_velocity; + float distance; + float n; +}; + +struct Env { + Log log; + Agent agents[1]; + int tag; + int boundary_reached; + int num_agents; + unsigned int rng; + const MjModel* m; + int tick; + float x_start; + float episode_return; + int max_steps; + float reset_noise_scale; + float forward_reward_weight; + float ctrl_cost_weight; + float contact_cost_weight; + float healthy_reward; + MjData d; +}; +typedef Env Humanoid; + +// Returns the contact cost (sum of squared external forces, clamped) +MJ_HD float compute_observations(Humanoid* env) { + const MjModel* m = env->m; + float* obs = env->agents[0].observations; + memcpy(obs, env->d.qpos + 2, (m->nq - 2)*sizeof(float)); + obs += m->nq - 2; + memcpy(obs, env->d.qvel, m->nv*sizeof(float)); + obs += m->nv; + memcpy(obs, env->d.cinert[1], 10*(m->nbody - 1)*sizeof(float)); + obs += 10*(m->nbody - 1); + memcpy(obs, env->d.cvel[1], 6*(m->nbody - 1)*sizeof(float)); + obs += 6*(m->nbody - 1); + memcpy(obs, env->d.qfrc_actuator + 6, (m->nv - 6)*sizeof(float)); + obs += m->nv - 6; + memcpy(obs, env->d.cfrc_ext[1], 6*(m->nbody - 1)*sizeof(float)); + float cost = 0.0f; + for (int b = 1; b < m->nbody; b++) { + cost += mju_dot6(env->d.cfrc_ext[b], env->d.cfrc_ext[b]); + } + return fminf(env->contact_cost_weight*cost, HM_CONTACT_COST_MAX); +} + +// x of the whole-body center of mass from the inertial frames of the last +// kinematics pass (Gym's mass_center) +MJ_HD float mjc_com_x(Humanoid* env) { + const MjModel* m = env->m; + float num = 0.0f; + float den = 0.0f; + for (int b = 0; b < m->nbody; b++) { + num += m->body_mass[b]*env->d.xipos[b][0]; + den += m->body_mass[b]; + } + return num / den; +} + +MJ_HD void mjc_reset(Humanoid* env) { + const MjModel* m = env->m; + mj_resetData(m, &env->d); + float s = env->reset_noise_scale; + for (int i = 0; i < m->nq; i++) { + env->d.qpos[i] += s*(2.0f*mju_rand(&env->rng) - 1.0f); + } + for (int i = 0; i < m->nv; i++) { + env->d.qvel[i] = s*(2.0f*mju_rand(&env->rng) - 1.0f); + } + mj_forward(m, &env->d); + env->tick = 0; + env->x_start = env->d.qpos[0]; + env->episode_return = 0.0f; + compute_observations(env); +} + +MJ_HD void mjc_step(Humanoid* env) { + const MjModel* m = env->m; + float* actions = env->agents[0].actions; + float cost = 0.0f; + for (int i = 0; i < m->nu; i++) { + const float* range = m->actuator_ctrlrange[i]; + env->d.ctrl[i] = fminf(fmaxf(actions[i], range[0]), range[1]); + cost += env->d.ctrl[i]*env->d.ctrl[i]; + } + float x0 = mjc_com_x(env); + for (int k = 0; k < HM_FRAME_SKIP; k++) { + mj_step(m, &env->d); + } + mj_rnePostConstraint(m, &env->d); + float* q = env->d.qpos; + float dt = HM_FRAME_SKIP*m->opt_timestep; + float x_velocity = (mjc_com_x(env) - x0) / dt; + int healthy = q[2] > 1.0f && q[2] < 2.0f; + float contact_cost = compute_observations(env); + float reward = env->forward_reward_weight*x_velocity - env->ctrl_cost_weight*cost + - contact_cost + (healthy ? env->healthy_reward : 0.0f); + env->tick++; + env->episode_return += reward; + env->agents[0].rewards[0] = reward; + env->agents[0].terminals[0] = 0.0f; + if (healthy && env->tick < env->max_steps) { + return; + } + float distance = q[0] - env->x_start; + float xvel = distance / (env->tick*dt); + env->agents[0].terminals[0] = 1.0f; + env->log.perf += fminf(fmaxf(xvel / 3.0f, 0.0f), 1.0f); + env->log.score += env->episode_return; + env->log.episode_return += env->episode_return; + env->log.episode_length += env->tick; + env->log.x_velocity += xvel; + env->log.distance += distance; + env->log.n += 1.0f; + mjc_reset(env); +} + +void mjc_render(Humanoid* env) { + float target[3] = {env->d.xpos[1][0], env->d.xpos[1][1], 1.0f}; + mj_render(env->m, &env->d, "PufferLib Humanoid", target, 5.0f, 1.5f, + TextFormat("step %d x %.2f m vel %.2f m/s return %.1f", env->tick, + env->d.qpos[0], env->d.qvel[0], env->episode_return)); +} + +void mjc_init(Humanoid* env, Dict* kwargs) { + if (mj_model.nbody == 0) { + mj_loadModel(&mj_model, dict_get_str(kwargs, "model")); + } + assert(mj_model.nq - 2 + 2*mj_model.nv - 6 + 22*(mj_model.nbody - 1) == OBS_SIZE); + assert(mj_model.nu == NUM_ATNS); + env->m = &mj_model; + env->num_agents = 1; + env->agents[0].policy = 0; + env->agents[0].action_mask = NULL; + env->max_steps = dict_get(kwargs, "max_steps"); + env->reset_noise_scale = dict_get(kwargs, "reset_noise_scale"); + env->forward_reward_weight = dict_get(kwargs, "forward_reward_weight"); + env->ctrl_cost_weight = dict_get(kwargs, "ctrl_cost_weight"); + env->contact_cost_weight = dict_get(kwargs, "contact_cost_weight"); + env->healthy_reward = dict_get(kwargs, "healthy_reward"); +} + +void puf_log(Log* log, Dict* out) { + dict_set(out, "perf", log->perf); + dict_set(out, "score", log->score); + dict_set(out, "episode_return", log->episode_return); + dict_set(out, "episode_length", log->episode_length); + dict_set(out, "x_velocity", log->x_velocity); + dict_set(out, "distance", log->distance); +} + +#include "backend.h" diff --git a/ocean/mujoco/mjc_walker2d.c b/ocean/mujoco/mjc_walker2d.c new file mode 100644 index 0000000000..771dd3ea3a --- /dev/null +++ b/ocean/mujoco/mjc_walker2d.c @@ -0,0 +1,3 @@ +#define PUFFERCPU_EVAL_MAIN +#define ENV_HEADER "../ocean/mujoco/mjc_walker2d.h" +#include "puffercpu.h" diff --git a/ocean/mujoco/mjc_walker2d.cu b/ocean/mujoco/mjc_walker2d.cu new file mode 100644 index 0000000000..c4f19d4494 --- /dev/null +++ b/ocean/mujoco/mjc_walker2d.cu @@ -0,0 +1,4 @@ +// GPU build of mjc_walker2d: one thread per env (see backend.h) +#define PUF_BACKEND PUF_GPU +typedef float obs_t; +#include "mjc_walker2d.h" diff --git a/ocean/mujoco/mjc_walker2d.h b/ocean/mujoco/mjc_walker2d.h new file mode 100644 index 0000000000..24386cc221 --- /dev/null +++ b/ocean/mujoco/mjc_walker2d.h @@ -0,0 +1,160 @@ +// Walker2d (gymnasium Walker2d-v5) on the MuJoCo-style physics core: obs = +// qpos[1:] + clip(qvel, -10, 10), reward = healthy + forward velocity - ctrl +// cost, terminates when the torso leaves z in (0.8, 2) or |angle| >= 1. +// Model: resources/mujoco/walker2d.xml (gymnasium walker2d_v5.xml) compiled by +// mjcf2bin.py. mjc_reset/mjc_step are host+device (see backend.h). + +#include +#include +#include +#include +#include "raylib.h" +typedef float obs_t; +#include "pufferenv.h" +// physics.h capacities sized to this model (checked by mj_loadModel) +#define MJ_MAX_NQ 9 +#define MJ_MAX_NV 9 +#define MJ_MAX_NBODY 8 +#define MJ_MAX_NJNT 9 +#define MJ_MAX_NGEOM 8 +#define MJ_MAX_NU 6 +#define MJ_MAXCON 14 +#define MJ_MAXEFC 62 +#include "physics.h" +#include "render.h" + +#define WK_FRAME_SKIP 4 +#define OBS_SIZE 17 +#define NUM_ATNS 6 +#define ACT_SIZES {1, 1, 1, 1, 1, 1} +#define PUF_STEPS_PER_SEC 125 + +MjModel mj_model; + +struct Log { + float perf; + float score; + float episode_return; + float episode_length; + float x_velocity; + float distance; + float n; +}; + +struct Env { + Log log; + Agent agents[1]; + int tag; + int boundary_reached; + int num_agents; + unsigned int rng; + const MjModel* m; + int tick; + float x_start; + float episode_return; + int max_steps; + float reset_noise_scale; + float forward_reward_weight; + float ctrl_cost_weight; + float healthy_reward; + MjData d; +}; +typedef Env Walker2d; + +MJ_HD void compute_observations(Walker2d* env) { + const MjModel* m = env->m; + float* obs = env->agents[0].observations; + memcpy(obs, env->d.qpos + 1, (m->nq - 1)*sizeof(float)); + for (int i = 0; i < m->nv; i++) { + obs[m->nq - 1 + i] = fminf(fmaxf(env->d.qvel[i], -10.0f), 10.0f); + } +} + +MJ_HD void mjc_reset(Walker2d* env) { + const MjModel* m = env->m; + mj_resetData(m, &env->d); + float s = env->reset_noise_scale; + for (int i = 0; i < m->nq; i++) { + env->d.qpos[i] += s*(2.0f*mju_rand(&env->rng) - 1.0f); + } + for (int i = 0; i < m->nv; i++) { + env->d.qvel[i] = s*(2.0f*mju_rand(&env->rng) - 1.0f); + } + env->tick = 0; + env->x_start = env->d.qpos[0]; + env->episode_return = 0.0f; + compute_observations(env); +} + +MJ_HD void mjc_step(Walker2d* env) { + const MjModel* m = env->m; + float* actions = env->agents[0].actions; + float cost = 0.0f; + for (int i = 0; i < m->nu; i++) { + env->d.ctrl[i] = fminf(fmaxf(actions[i], -1.0f), 1.0f); + cost += env->d.ctrl[i]*env->d.ctrl[i]; + } + float x0 = env->d.qpos[0]; + for (int k = 0; k < WK_FRAME_SKIP; k++) { + mj_step(m, &env->d); + } + float* q = env->d.qpos; + float dt = WK_FRAME_SKIP*m->opt_timestep; + float x_velocity = (q[0] - x0) / dt; + int healthy = q[1] > 0.8f && q[1] < 2.0f && q[2] > -1.0f && q[2] < 1.0f; + float reward = env->forward_reward_weight*x_velocity - env->ctrl_cost_weight*cost + + (healthy ? env->healthy_reward : 0.0f); + env->tick++; + env->episode_return += reward; + env->agents[0].rewards[0] = reward; + env->agents[0].terminals[0] = 0.0f; + if (healthy && env->tick < env->max_steps) { + compute_observations(env); + return; + } + float distance = q[0] - env->x_start; + float xvel = distance / (env->tick*dt); + env->agents[0].terminals[0] = 1.0f; + env->log.perf += fminf(fmaxf(xvel / 5.0f, 0.0f), 1.0f); + env->log.score += env->episode_return; + env->log.episode_return += env->episode_return; + env->log.episode_length += env->tick; + env->log.x_velocity += xvel; + env->log.distance += distance; + env->log.n += 1.0f; + mjc_reset(env); +} + +void mjc_render(Walker2d* env) { + float target[3] = {env->d.xpos[1][0], 0.0f, 0.8f}; + mj_render(env->m, &env->d, "PufferLib Walker2d", target, 4.0f, 1.0f, + TextFormat("step %d x %.2f m vel %.2f m/s return %.1f", env->tick, + env->d.qpos[0], env->d.qvel[0], env->episode_return)); +} + +void mjc_init(Walker2d* env, Dict* kwargs) { + if (mj_model.nbody == 0) { + mj_loadModel(&mj_model, dict_get_str(kwargs, "model")); + } + assert(mj_model.nq - 1 + mj_model.nv == OBS_SIZE && mj_model.nu == NUM_ATNS); + env->m = &mj_model; + env->num_agents = 1; + env->agents[0].policy = 0; + env->agents[0].action_mask = NULL; + env->max_steps = dict_get(kwargs, "max_steps"); + env->reset_noise_scale = dict_get(kwargs, "reset_noise_scale"); + env->forward_reward_weight = dict_get(kwargs, "forward_reward_weight"); + env->ctrl_cost_weight = dict_get(kwargs, "ctrl_cost_weight"); + env->healthy_reward = dict_get(kwargs, "healthy_reward"); +} + +void puf_log(Log* log, Dict* out) { + dict_set(out, "perf", log->perf); + dict_set(out, "score", log->score); + dict_set(out, "episode_return", log->episode_return); + dict_set(out, "episode_length", log->episode_length); + dict_set(out, "x_velocity", log->x_velocity); + dict_set(out, "distance", log->distance); +} + +#include "backend.h" diff --git a/ocean/mujoco/mjcf2bin.py b/ocean/mujoco/mjcf2bin.py new file mode 100644 index 0000000000..84ef892cf3 --- /dev/null +++ b/ocean/mujoco/mjcf2bin.py @@ -0,0 +1,73 @@ +"""Compile MJCF models with MuJoCo into the binaries loaded by physics.h. + + python ocean/mujoco/mjcf2bin.py # every resources/mujoco/*.xml + python ocean/mujoco/mjcf2bin.py path/to/model.xml [out.bin] + +Each .bin is a header (magic, version, nq, nv, nbody, njnt, ngeom, nsite, +nu) followed by the compiled mjModel arrays in the order read by mj_loadModel, +as int32/float32. Compiling with MuJoCo keeps inertias, invweights and bounding +radii bit-identical to the real thing. Requires pip install mujoco. +""" +import glob +import os +import struct +import sys + +import mujoco +import numpy as np + + +def compile_model(xml, out): + m = mujoco.MjModel.from_xml_path(xml) + assert m.neq == 0, "equality constraints unsupported" + assert all(m.actuator_trntype == 0) and all(m.actuator_dyntype == 0), "motors only" + assert all(m.jnt_type[m.actuator_trnid[:, 0]] >= 2), "actuators on hinge/slide joints only" + assert m.opt.cone == 0, "pyramidal cones only" + # fixed tendons without limits, springs or dampers (humanoid) exert no force + assert not m.tendon_limited.any() and not m.tendon_stiffness.any() \ + and not m.tendon_damping.any(), "tendon limits/springs/dampers unsupported" + i32, f32 = np.int32, np.float32 + fields = [ + (m.opt.timestep, f32), (m.opt.gravity, f32), (m.opt.integrator, i32), + (m.opt.impratio, f32), + (m.body_parentid, i32), (m.body_rootid, i32), (m.body_weldid, i32), + (m.body_jntadr, i32), (m.body_jntnum, i32), (m.body_dofadr, i32), (m.body_dofnum, i32), + (m.body_pos, f32), (m.body_quat, f32), (m.body_ipos, f32), (m.body_iquat, f32), + (m.body_mass, f32), (m.body_subtreemass, f32), (m.body_inertia, f32), + (m.body_invweight0, f32), + (m.jnt_type, i32), (m.jnt_qposadr, i32), (m.jnt_dofadr, i32), (m.jnt_limited, i32), + (m.jnt_axis, f32), (m.jnt_pos, f32), (m.jnt_range, f32), (m.jnt_margin, f32), + (m.jnt_stiffness, f32), (m.jnt_solref, f32), (m.jnt_solimp, f32), + (m.dof_bodyid, i32), (m.dof_jntid, i32), (m.dof_parentid, i32), (m.dof_armature, f32), + (m.dof_damping, f32), (m.dof_invweight0, f32), + (m.geom_type, i32), (m.geom_bodyid, i32), (m.geom_contype, i32), + (m.geom_conaffinity, i32), (m.geom_condim, i32), (m.geom_priority, i32), + (m.geom_solmix, f32), (m.geom_size, f32), (m.geom_pos, f32), (m.geom_quat, f32), + (m.geom_friction, f32), (m.geom_solref, f32), (m.geom_solimp, f32), + (m.geom_margin, f32), (m.geom_gap, f32), (m.geom_rbound, f32), + (m.site_bodyid, i32), (m.site_pos, f32), (m.site_quat, f32), + (m.actuator_trnid[:, 0], i32), (m.actuator_gear[:, 0], f32), + (m.actuator_ctrllimited, i32), (m.actuator_ctrlrange, f32), + (m.qpos0, f32), (m.qpos_spring, f32), + ] + with open(out, "wb") as f: + f.write(struct.pack("9i", 0x4E424A4D, 1, m.nq, m.nv, m.nbody, m.njnt, m.ngeom, m.nsite, + m.nu)) + for value, dtype in fields: + f.write(np.ascontiguousarray(value, dtype=dtype).tobytes()) + print("wrote %s: nq %d nv %d nbody %d njnt %d ngeom %d nsite %d nu %d" % (out, m.nq, m.nv, + m.nbody, m.njnt, m.ngeom, m.nsite, m.nu)) + + +def main(): + if len(sys.argv) > 1: + xml = sys.argv[1] + compile_model(xml, sys.argv[2] if len(sys.argv) > 2 else os.path.splitext(xml)[0] + ".bin") + return + root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + for xml in sorted(glob.glob(os.path.join(root, "resources", "mujoco", "*.xml"))): + compile_model(xml, os.path.splitext(xml)[0] + ".bin") + + +if __name__ == "__main__": + main() diff --git a/ocean/mujoco/physics.h b/ocean/mujoco/physics.h new file mode 100644 index 0000000000..addb9c13e9 --- /dev/null +++ b/ocean/mujoco/physics.h @@ -0,0 +1,1484 @@ +// MuJoCo-style rigid body physics for PufferLib envs. CUDA C99, fixed capacity, +// float. Ports mj_step's pipeline (same algorithms, names and layouts as the +// MuJoCo engine) for the features the classic Gym models use: free/ball/hinge/ +// slide joints, plane/sphere/capsule geoms, motor actuators, joint springs and +// dampers, joint limits and pyramidal friction contacts with MuJoCo's soft +// constraint model, semi-implicit Euler with implicit damping, and RK4. +// Models are compiled from MJCF by ocean/mujoco/mjcf2bin.py and loaded with +// mj_loadModel; arrays have fixed MJ_MAX_* capacity (defaults below, envs +// define tighter ones for their model before including) and runtime counts. + +#include +#include +#include +#include +#include +#include +#ifdef __CUDACC__ +#define MJ_HD __host__ __device__ +#else +#define MJ_HD +#endif + +#define MJ_MINVAL 1e-15f +#define MJ_MINIMP 0.0001f +#define MJ_MAXIMP 0.9999f +#ifndef MJ_MAX_NQ +#define MJ_MAX_NQ 32 +#endif +#ifndef MJ_MAX_NV +#define MJ_MAX_NV 32 +#endif +#ifndef MJ_MAX_NBODY +#define MJ_MAX_NBODY 16 +#endif +#ifndef MJ_MAX_NJNT +#define MJ_MAX_NJNT 24 +#endif +#ifndef MJ_MAX_NGEOM +#define MJ_MAX_NGEOM 24 +#endif +#ifndef MJ_MAX_NSITE +#define MJ_MAX_NSITE 4 +#endif +#ifndef MJ_MAX_NU +#define MJ_MAX_NU 24 +#endif +#ifndef MJ_MAXCON +#define MJ_MAXCON 48 +#endif +// constraint rows: 1 per active limit, 1 per frictionless contact, 2*(condim-1) otherwise +#ifndef MJ_MAXEFC +#define MJ_MAXEFC 128 +#endif +#define MJ_MAGIC 0x4e424a4d +enum {MJ_JNT_FREE, MJ_JNT_BALL, MJ_JNT_SLIDE, MJ_JNT_HINGE}; +enum {MJ_GEOM_PLANE, MJ_GEOM_HFIELD, MJ_GEOM_SPHERE, MJ_GEOM_CAPSULE, MJ_GEOM_ELLIPSOID, + MJ_GEOM_CYLINDER, MJ_GEOM_BOX}; +enum {MJ_INT_EULER, MJ_INT_RK4}; + +typedef struct { + int nq, nv, nbody, njnt, ngeom, nsite, nu; + float opt_timestep; + float opt_gravity[3]; + int opt_integrator; + float opt_impratio; + int body_parentid[MJ_MAX_NBODY]; + int body_rootid[MJ_MAX_NBODY]; + int body_weldid[MJ_MAX_NBODY]; + int body_jntadr[MJ_MAX_NBODY]; + int body_jntnum[MJ_MAX_NBODY]; + int body_dofadr[MJ_MAX_NBODY]; + int body_dofnum[MJ_MAX_NBODY]; + float body_pos[MJ_MAX_NBODY][3]; + float body_quat[MJ_MAX_NBODY][4]; + float body_ipos[MJ_MAX_NBODY][3]; + float body_iquat[MJ_MAX_NBODY][4]; + float body_mass[MJ_MAX_NBODY]; + float body_subtreemass[MJ_MAX_NBODY]; + float body_inertia[MJ_MAX_NBODY][3]; + float body_invweight0[MJ_MAX_NBODY][2]; + int jnt_type[MJ_MAX_NJNT]; + int jnt_qposadr[MJ_MAX_NJNT]; + int jnt_dofadr[MJ_MAX_NJNT]; + int jnt_limited[MJ_MAX_NJNT]; + float jnt_axis[MJ_MAX_NJNT][3]; + float jnt_pos[MJ_MAX_NJNT][3]; + float jnt_range[MJ_MAX_NJNT][2]; + float jnt_margin[MJ_MAX_NJNT]; + float jnt_stiffness[MJ_MAX_NJNT]; + float jnt_solref[MJ_MAX_NJNT][2]; + float jnt_solimp[MJ_MAX_NJNT][5]; + int dof_bodyid[MJ_MAX_NV]; + int dof_jntid[MJ_MAX_NV]; + int dof_parentid[MJ_MAX_NV]; + float dof_armature[MJ_MAX_NV]; + float dof_damping[MJ_MAX_NV]; + float dof_invweight0[MJ_MAX_NV]; + int geom_type[MJ_MAX_NGEOM]; + int geom_bodyid[MJ_MAX_NGEOM]; + int geom_contype[MJ_MAX_NGEOM]; + int geom_conaffinity[MJ_MAX_NGEOM]; + int geom_condim[MJ_MAX_NGEOM]; + int geom_priority[MJ_MAX_NGEOM]; + float geom_solmix[MJ_MAX_NGEOM]; + float geom_size[MJ_MAX_NGEOM][3]; + float geom_pos[MJ_MAX_NGEOM][3]; + float geom_quat[MJ_MAX_NGEOM][4]; + float geom_friction[MJ_MAX_NGEOM][3]; + float geom_solref[MJ_MAX_NGEOM][2]; + float geom_solimp[MJ_MAX_NGEOM][5]; + float geom_margin[MJ_MAX_NGEOM]; + float geom_gap[MJ_MAX_NGEOM]; + float geom_rbound[MJ_MAX_NGEOM]; + int site_bodyid[MJ_MAX_NSITE]; + float site_pos[MJ_MAX_NSITE][3]; + float site_quat[MJ_MAX_NSITE][4]; + int actuator_trnid[MJ_MAX_NU]; + float actuator_gear[MJ_MAX_NU]; + int actuator_ctrllimited[MJ_MAX_NU]; + float actuator_ctrlrange[MJ_MAX_NU][2]; + float qpos0[MJ_MAX_NQ]; + float qpos_spring[MJ_MAX_NQ]; +} MjModel; + +typedef struct { + int geom1, geom2, dim; + float dist; + float pos[3]; + float frame[9]; + float includemargin; + float friction[5]; + float solref[2]; + float solimp[5]; + int efc_address; +} MjContact; + +// Constraint solver scratch: MJ_SCRATCH floats per env +#define MJ_SCRATCH (2*MJ_MAXEFC*MJ_MAXEFC + MJ_MAXEFC*MJ_MAX_NV) + +typedef struct { + float* efc_AR; + float* efc_ARfree; + float* efc_MinvJT; + int efc_stride; + float time; + float qpos[MJ_MAX_NQ]; + float qvel[MJ_MAX_NV]; + float ctrl[MJ_MAX_NU]; + float qacc[MJ_MAX_NV]; + float xpos[MJ_MAX_NBODY][3]; + float xipos[MJ_MAX_NBODY][3]; + float xquat[MJ_MAX_NBODY][4]; + float xmat[MJ_MAX_NBODY][9]; + float ximat[MJ_MAX_NBODY][9]; + float xanchor[MJ_MAX_NJNT][3]; + float xaxis[MJ_MAX_NJNT][3]; + float geom_xpos[MJ_MAX_NGEOM][3]; + float geom_xmat[MJ_MAX_NGEOM][9]; + float site_xpos[MJ_MAX_NSITE][3]; + float site_xmat[MJ_MAX_NSITE][9]; + float subtree_com[MJ_MAX_NBODY][3]; + float cinert[MJ_MAX_NBODY][10]; + float cdof[MJ_MAX_NV][6]; + float cvel[MJ_MAX_NBODY][6]; + float cdof_dot[MJ_MAX_NV][6]; + float qM[MJ_MAX_NV][MJ_MAX_NV]; + float qLD[MJ_MAX_NV][MJ_MAX_NV]; + float qfrc_bias[MJ_MAX_NV]; + float qfrc_passive[MJ_MAX_NV]; + float qfrc_actuator[MJ_MAX_NV]; + float qfrc_smooth[MJ_MAX_NV]; + float qacc_smooth[MJ_MAX_NV]; + float qfrc_constraint[MJ_MAX_NV]; + float cacc[MJ_MAX_NBODY][6]; + float cfrc_int[MJ_MAX_NBODY][6]; + float cfrc_ext[MJ_MAX_NBODY][6]; + int ncon; + MjContact contact[MJ_MAXCON]; + int nefc; + float efc_J[MJ_MAXEFC][MJ_MAX_NV]; + float efc_pos[MJ_MAXEFC]; + float efc_margin[MJ_MAXEFC]; + float efc_R[MJ_MAXEFC]; + float efc_aref[MJ_MAXEFC]; + float efc_force[MJ_MAXEFC]; +} MjData; + +MJ_HD void mj_makeData(MjData* d, float* scratch, int stride) { + d->efc_AR = scratch; + d->efc_ARfree = scratch + MJ_MAXEFC*MJ_MAXEFC*stride; + d->efc_MinvJT = scratch + 2*MJ_MAXEFC*MJ_MAXEFC*stride; + d->efc_stride = stride; +} + +void mj_read(FILE* fp, void* dst, int count, int size) { + assert((int)fread(dst, size, count, fp) == count && "truncated model file"); +} + +// Load a model compiled by ocean/mujoco/mjcf2bin.py (same field order) +void mj_loadModel(MjModel* m, const char* path) { + FILE* fp = fopen(path, "rb"); + assert(fp && "cannot open model file"); + int head[9]; + mj_read(fp, head, 9, sizeof(int)); + assert(head[0] == MJ_MAGIC && head[1] == 1 && "bad model file"); + m->nq = head[2]; + m->nv = head[3]; + m->nbody = head[4]; + m->njnt = head[5]; + m->ngeom = head[6]; + m->nsite = head[7]; + m->nu = head[8]; + assert(m->nq <= MJ_MAX_NQ && m->nv <= MJ_MAX_NV && m->nbody <= MJ_MAX_NBODY + && m->njnt <= MJ_MAX_NJNT && m->ngeom <= MJ_MAX_NGEOM && m->nsite <= MJ_MAX_NSITE + && m->nu <= MJ_MAX_NU && "model exceeds MJ_MAX_* capacity"); + mj_read(fp, &m->opt_timestep, 1, sizeof(float)); + mj_read(fp, m->opt_gravity, 3, sizeof(float)); + mj_read(fp, &m->opt_integrator, 1, sizeof(int)); + mj_read(fp, &m->opt_impratio, 1, sizeof(float)); + mj_read(fp, m->body_parentid, m->nbody, sizeof(int)); + mj_read(fp, m->body_rootid, m->nbody, sizeof(int)); + mj_read(fp, m->body_weldid, m->nbody, sizeof(int)); + mj_read(fp, m->body_jntadr, m->nbody, sizeof(int)); + mj_read(fp, m->body_jntnum, m->nbody, sizeof(int)); + mj_read(fp, m->body_dofadr, m->nbody, sizeof(int)); + mj_read(fp, m->body_dofnum, m->nbody, sizeof(int)); + mj_read(fp, m->body_pos, 3*m->nbody, sizeof(float)); + mj_read(fp, m->body_quat, 4*m->nbody, sizeof(float)); + mj_read(fp, m->body_ipos, 3*m->nbody, sizeof(float)); + mj_read(fp, m->body_iquat, 4*m->nbody, sizeof(float)); + mj_read(fp, m->body_mass, m->nbody, sizeof(float)); + mj_read(fp, m->body_subtreemass, m->nbody, sizeof(float)); + mj_read(fp, m->body_inertia, 3*m->nbody, sizeof(float)); + mj_read(fp, m->body_invweight0, 2*m->nbody, sizeof(float)); + mj_read(fp, m->jnt_type, m->njnt, sizeof(int)); + mj_read(fp, m->jnt_qposadr, m->njnt, sizeof(int)); + mj_read(fp, m->jnt_dofadr, m->njnt, sizeof(int)); + mj_read(fp, m->jnt_limited, m->njnt, sizeof(int)); + mj_read(fp, m->jnt_axis, 3*m->njnt, sizeof(float)); + mj_read(fp, m->jnt_pos, 3*m->njnt, sizeof(float)); + mj_read(fp, m->jnt_range, 2*m->njnt, sizeof(float)); + mj_read(fp, m->jnt_margin, m->njnt, sizeof(float)); + mj_read(fp, m->jnt_stiffness, m->njnt, sizeof(float)); + mj_read(fp, m->jnt_solref, 2*m->njnt, sizeof(float)); + mj_read(fp, m->jnt_solimp, 5*m->njnt, sizeof(float)); + mj_read(fp, m->dof_bodyid, m->nv, sizeof(int)); + mj_read(fp, m->dof_jntid, m->nv, sizeof(int)); + mj_read(fp, m->dof_parentid, m->nv, sizeof(int)); + mj_read(fp, m->dof_armature, m->nv, sizeof(float)); + mj_read(fp, m->dof_damping, m->nv, sizeof(float)); + mj_read(fp, m->dof_invweight0, m->nv, sizeof(float)); + mj_read(fp, m->geom_type, m->ngeom, sizeof(int)); + mj_read(fp, m->geom_bodyid, m->ngeom, sizeof(int)); + mj_read(fp, m->geom_contype, m->ngeom, sizeof(int)); + mj_read(fp, m->geom_conaffinity, m->ngeom, sizeof(int)); + mj_read(fp, m->geom_condim, m->ngeom, sizeof(int)); + mj_read(fp, m->geom_priority, m->ngeom, sizeof(int)); + mj_read(fp, m->geom_solmix, m->ngeom, sizeof(float)); + mj_read(fp, m->geom_size, 3*m->ngeom, sizeof(float)); + mj_read(fp, m->geom_pos, 3*m->ngeom, sizeof(float)); + mj_read(fp, m->geom_quat, 4*m->ngeom, sizeof(float)); + mj_read(fp, m->geom_friction, 3*m->ngeom, sizeof(float)); + mj_read(fp, m->geom_solref, 2*m->ngeom, sizeof(float)); + mj_read(fp, m->geom_solimp, 5*m->ngeom, sizeof(float)); + mj_read(fp, m->geom_margin, m->ngeom, sizeof(float)); + mj_read(fp, m->geom_gap, m->ngeom, sizeof(float)); + mj_read(fp, m->geom_rbound, m->ngeom, sizeof(float)); + mj_read(fp, m->site_bodyid, m->nsite, sizeof(int)); + mj_read(fp, m->site_pos, 3*m->nsite, sizeof(float)); + mj_read(fp, m->site_quat, 4*m->nsite, sizeof(float)); + mj_read(fp, m->actuator_trnid, m->nu, sizeof(int)); + mj_read(fp, m->actuator_gear, m->nu, sizeof(float)); + mj_read(fp, m->actuator_ctrllimited, m->nu, sizeof(int)); + mj_read(fp, m->actuator_ctrlrange, 2*m->nu, sizeof(float)); + mj_read(fp, m->qpos0, m->nq, sizeof(float)); + mj_read(fp, m->qpos_spring, m->nq, sizeof(float)); + fclose(fp); +} + +// Vector, quaternion and spatial algebra (mju_*) + +MJ_HD float mju_dot3(const float* a, const float* b) { + return a[0]*b[0] + a[1]*b[1] + a[2]*b[2]; +} + +MJ_HD void mju_cross(float* res, const float* a, const float* b) { + float t0 = a[1]*b[2] - a[2]*b[1]; + float t1 = a[2]*b[0] - a[0]*b[2]; + float t2 = a[0]*b[1] - a[1]*b[0]; + res[0] = t0; + res[1] = t1; + res[2] = t2; +} + +MJ_HD float mju_normalize3(float* v) { + float norm = sqrtf(mju_dot3(v, v)); + if (norm < MJ_MINVAL) { + v[0] = 1.0f; + v[1] = 0.0f; + v[2] = 0.0f; + } else { + v[0] /= norm; + v[1] /= norm; + v[2] /= norm; + } + return norm; +} + +MJ_HD void mju_normalize4(float* q) { + float norm = sqrtf(q[0]*q[0] + q[1]*q[1] + q[2]*q[2] + q[3]*q[3]); + if (norm < MJ_MINVAL) { + q[0] = 1.0f; + q[1] = 0.0f; + q[2] = 0.0f; + q[3] = 0.0f; + } else { + q[0] /= norm; + q[1] /= norm; + q[2] /= norm; + q[3] /= norm; + } +} + +MJ_HD void mju_mulQuat(float* res, const float* qa, const float* qb) { + res[0] = qa[0]*qb[0] - qa[1]*qb[1] - qa[2]*qb[2] - qa[3]*qb[3]; + res[1] = qa[0]*qb[1] + qa[1]*qb[0] + qa[2]*qb[3] - qa[3]*qb[2]; + res[2] = qa[0]*qb[2] - qa[1]*qb[3] + qa[2]*qb[0] + qa[3]*qb[1]; + res[3] = qa[0]*qb[3] + qa[1]*qb[2] - qa[2]*qb[1] + qa[3]*qb[0]; +} + +MJ_HD void mju_rotVecQuat(float* res, const float* vec, const float* quat) { + float t0 = quat[0]*vec[0] + quat[2]*vec[2] - quat[3]*vec[1]; + float t1 = quat[0]*vec[1] + quat[3]*vec[0] - quat[1]*vec[2]; + float t2 = quat[0]*vec[2] + quat[1]*vec[1] - quat[2]*vec[0]; + res[0] = vec[0] + 2.0f*(quat[2]*t2 - quat[3]*t1); + res[1] = vec[1] + 2.0f*(quat[3]*t0 - quat[1]*t2); + res[2] = vec[2] + 2.0f*(quat[1]*t1 - quat[2]*t0); +} + +MJ_HD void mju_quat2Mat(float* res, const float* q) { + float q00 = q[0]*q[0], q01 = q[0]*q[1], q02 = q[0]*q[2], q03 = q[0]*q[3]; + float q11 = q[1]*q[1], q12 = q[1]*q[2], q13 = q[1]*q[3]; + float q22 = q[2]*q[2], q23 = q[2]*q[3], q33 = q[3]*q[3]; + res[0] = q00 + q11 - q22 - q33; + res[4] = q00 - q11 + q22 - q33; + res[8] = q00 - q11 - q22 + q33; + res[1] = 2.0f*(q12 - q03); + res[2] = 2.0f*(q13 + q02); + res[3] = 2.0f*(q12 + q03); + res[5] = 2.0f*(q23 - q01); + res[6] = 2.0f*(q13 - q02); + res[7] = 2.0f*(q23 + q01); +} + +MJ_HD void mju_axisAngle2Quat(float* res, const float* axis, float angle) { + float s = sinf(0.5f*angle); + res[0] = cosf(0.5f*angle); + res[1] = axis[0]*s; + res[2] = axis[1]*s; + res[3] = axis[2]*s; +} + +// Integrate quaternion by angular velocity expressed in the local frame +MJ_HD void mju_quatIntegrate(float* quat, const float* vel, float scale) { + float axis[3] = {vel[0], vel[1], vel[2]}; + float angle = scale*mju_normalize3(axis); + float qrot[4]; + mju_axisAngle2Quat(qrot, axis, angle); + mju_normalize4(quat); + mju_mulQuat(quat, quat, qrot); +} + +MJ_HD void mju_mulMatTVec3(float* res, const float* mat, const float* vec) { + res[0] = mat[0]*vec[0] + mat[3]*vec[1] + mat[6]*vec[2]; + res[1] = mat[1]*vec[0] + mat[4]*vec[1] + mat[7]*vec[2]; + res[2] = mat[2]*vec[0] + mat[5]*vec[1] + mat[8]*vec[2]; +} + +MJ_HD void mju_mulMatVec3(float* res, const float* mat, const float* vec) { + res[0] = mat[0]*vec[0] + mat[1]*vec[1] + mat[2]*vec[2]; + res[1] = mat[3]*vec[0] + mat[4]*vec[1] + mat[5]*vec[2]; + res[2] = mat[6]*vec[0] + mat[7]*vec[1] + mat[8]*vec[2]; +} + +// Motion vectors are (angular, linear); cross products of motion and force +MJ_HD void mju_crossMotion(float* res, const float* vel, const float* v) { + res[0] = -vel[2]*v[1] + vel[1]*v[2]; + res[1] = vel[2]*v[0] - vel[0]*v[2]; + res[2] = -vel[1]*v[0] + vel[0]*v[1]; + res[3] = -vel[2]*v[4] + vel[1]*v[5] - vel[5]*v[1] + vel[4]*v[2]; + res[4] = vel[2]*v[3] - vel[0]*v[5] + vel[5]*v[0] - vel[3]*v[2]; + res[5] = -vel[1]*v[3] + vel[0]*v[4] - vel[4]*v[0] + vel[3]*v[1]; +} + +MJ_HD void mju_crossForce(float* res, const float* vel, const float* f) { + res[0] = -vel[2]*f[1] + vel[1]*f[2] - vel[5]*f[4] + vel[4]*f[5]; + res[1] = vel[2]*f[0] - vel[0]*f[2] + vel[5]*f[3] - vel[3]*f[5]; + res[2] = -vel[1]*f[0] + vel[0]*f[1] - vel[4]*f[3] + vel[3]*f[4]; + res[3] = -vel[2]*f[4] + vel[1]*f[5]; + res[4] = vel[2]*f[3] - vel[0]*f[5]; + res[5] = -vel[1]*f[3] + vel[0]*f[4]; +} + +// Spatial inertia (10: I_xx I_yy I_zz I_xy I_xz I_yz, m*com, m) in a frame +// displaced by dif from the body inertial frame, rotated by mat +MJ_HD void mju_inertCom(float* res, const float* inert, const float* mat, const float* dif, + float mass) { + float tmp[9] = {mat[0]*inert[0], mat[3]*inert[0], mat[6]*inert[0], + mat[1]*inert[1], mat[4]*inert[1], mat[7]*inert[1], + mat[2]*inert[2], mat[5]*inert[2], mat[8]*inert[2]}; + res[0] = mat[0]*tmp[0] + mat[1]*tmp[3] + mat[2]*tmp[6] + mass*(dif[1]*dif[1] + dif[2]*dif[2]); + res[1] = mat[3]*tmp[1] + mat[4]*tmp[4] + mat[5]*tmp[7] + mass*(dif[0]*dif[0] + dif[2]*dif[2]); + res[2] = mat[6]*tmp[2] + mat[7]*tmp[5] + mat[8]*tmp[8] + mass*(dif[0]*dif[0] + dif[1]*dif[1]); + res[3] = mat[0]*tmp[1] + mat[1]*tmp[4] + mat[2]*tmp[7] - mass*dif[0]*dif[1]; + res[4] = mat[0]*tmp[2] + mat[1]*tmp[5] + mat[2]*tmp[8] - mass*dif[0]*dif[2]; + res[5] = mat[3]*tmp[2] + mat[4]*tmp[5] + mat[5]*tmp[8] - mass*dif[1]*dif[2]; + res[6] = mass*dif[0]; + res[7] = mass*dif[1]; + res[8] = mass*dif[2]; + res[9] = mass; +} + +MJ_HD void mju_mulInertVec(float* res, const float* i, const float* v) { + res[0] = i[0]*v[0] + i[3]*v[1] + i[4]*v[2] - i[8]*v[4] + i[7]*v[5]; + res[1] = i[3]*v[0] + i[1]*v[1] + i[5]*v[2] + i[8]*v[3] - i[6]*v[5]; + res[2] = i[4]*v[0] + i[5]*v[1] + i[2]*v[2] - i[7]*v[3] + i[6]*v[4]; + res[3] = i[8]*v[1] - i[7]*v[2] + i[9]*v[3]; + res[4] = i[6]*v[2] - i[8]*v[0] + i[9]*v[4]; + res[5] = i[7]*v[0] - i[6]*v[1] + i[9]*v[5]; +} + +MJ_HD float mju_dot6(const float* a, const float* b) { + return a[0]*b[0] + a[1]*b[1] + a[2]*b[2] + a[3]*b[3] + a[4]*b[4] + a[5]*b[5]; +} + +// xorshift32 uniform in [0, 1) and standard normal, for env resets +MJ_HD float mju_rand(unsigned int* rng) { + unsigned int x = *rng ? *rng : 0x9e3779b9u; + x ^= x << 13; + x ^= x >> 17; + x ^= x << 5; + *rng = x; + return (x >> 8)*(1.0f / 16777216.0f); +} + +MJ_HD float mju_randn(unsigned int* rng) { + float u1 = 1.0f - mju_rand(rng); + float u2 = mju_rand(rng); + return sqrtf(-2.0f*logf(u1))*cosf(2.0f*(float)M_PI*u2); +} + +// In-place Cholesky factor A = L L^T +MJ_HD void mju_cholFactor(float* A, int n, int lda, int es) { + for (int j = 0; j < n; j++) { + float d = A[(j*lda + j)*es]; + for (int k = 0; k < j; k++) { + d -= A[(j*lda + k)*es]*A[(j*lda + k)*es]; + } + d = sqrtf(d); + A[(j*lda + j)*es] = d; + for (int i = j + 1; i < n; i++) { + float s = A[(i*lda + j)*es]; + for (int k = 0; k < j; k++) { + s -= A[(i*lda + k)*es]*A[(j*lda + k)*es]; + } + A[(i*lda + j)*es] = s / d; + } + } +} + +MJ_HD void mju_cholSolve(const float* L, int n, int lda, int es, float* x) { + for (int i = 0; i < n; i++) { + float s = x[i]; + for (int k = 0; k < i; k++) { + s -= L[(i*lda + k)*es]*x[k]; + } + x[i] = s / L[(i*lda + i)*es]; + } + for (int i = n - 1; i >= 0; i--) { + float s = x[i]; + for (int k = i + 1; k < n; k++) { + s -= L[(k*lda + i)*es]*x[k]; + } + x[i] = s / L[(i*lda + i)*es]; + } +} + +// Forward kinematics: body, inertial, geom and site frames, joint anchors/axes + +MJ_HD void mj_local2Global(MjData* d, float* xpos, float* xmat, const float* pos, + const float* quat, int body) { + float tmp[3], q[4]; + mju_mulMatVec3(tmp, d->xmat[body], pos); + xpos[0] = d->xpos[body][0] + tmp[0]; + xpos[1] = d->xpos[body][1] + tmp[1]; + xpos[2] = d->xpos[body][2] + tmp[2]; + mju_mulQuat(q, d->xquat[body], quat); + mju_quat2Mat(xmat, q); +} + +MJ_HD void mj_kinematics(const MjModel* m, MjData* d) { + memset(d->xpos[0], 0, 3*sizeof(float)); + memset(d->xmat[0], 0, 9*sizeof(float)); + d->xquat[0][0] = 1.0f; + d->xquat[0][1] = 0.0f; + d->xquat[0][2] = 0.0f; + d->xquat[0][3] = 0.0f; + d->xmat[0][0] = 1.0f; + d->xmat[0][4] = 1.0f; + d->xmat[0][8] = 1.0f; + for (int i = 1; i < m->nbody; i++) { + float xpos[3], xquat[4]; + int jntadr = m->body_jntadr[i]; + int jntnum = m->body_jntnum[i]; + if (jntnum == 1 && m->jnt_type[jntadr] == MJ_JNT_FREE) { + int qadr = m->jnt_qposadr[jntadr]; + memcpy(xpos, d->qpos + qadr, 3*sizeof(float)); + memcpy(xquat, d->qpos + qadr + 3, 4*sizeof(float)); + mju_normalize4(xquat); + memcpy(d->xanchor[jntadr], xpos, 3*sizeof(float)); + memcpy(d->xaxis[jntadr], m->jnt_axis[jntadr], 3*sizeof(float)); + } else { + int pid = m->body_parentid[i]; + mju_mulMatVec3(xpos, d->xmat[pid], m->body_pos[i]); + xpos[0] += d->xpos[pid][0]; + xpos[1] += d->xpos[pid][1]; + xpos[2] += d->xpos[pid][2]; + mju_mulQuat(xquat, d->xquat[pid], m->body_quat[i]); + for (int j = jntadr; j < jntadr + jntnum; j++) { + int qadr = m->jnt_qposadr[j]; + int jtype = m->jnt_type[j]; + float xaxis[3], xanchor[3]; + mju_rotVecQuat(xaxis, m->jnt_axis[j], xquat); + mju_rotVecQuat(xanchor, m->jnt_pos[j], xquat); + xanchor[0] += xpos[0]; + xanchor[1] += xpos[1]; + xanchor[2] += xpos[2]; + if (jtype == MJ_JNT_SLIDE) { + float dq = d->qpos[qadr] - m->qpos0[qadr]; + xpos[0] += xaxis[0]*dq; + xpos[1] += xaxis[1]*dq; + xpos[2] += xaxis[2]*dq; + } else { + float qloc[4], vec[3]; + if (jtype == MJ_JNT_BALL) { + memcpy(qloc, d->qpos + qadr, 4*sizeof(float)); + mju_normalize4(qloc); + } else { + mju_axisAngle2Quat(qloc, m->jnt_axis[j], d->qpos[qadr] - m->qpos0[qadr]); + } + mju_mulQuat(xquat, xquat, qloc); + mju_rotVecQuat(vec, m->jnt_pos[j], xquat); + xpos[0] = xanchor[0] - vec[0]; + xpos[1] = xanchor[1] - vec[1]; + xpos[2] = xanchor[2] - vec[2]; + } + memcpy(d->xanchor[j], xanchor, 3*sizeof(float)); + memcpy(d->xaxis[j], xaxis, 3*sizeof(float)); + } + } + mju_normalize4(xquat); + memcpy(d->xpos[i], xpos, 3*sizeof(float)); + memcpy(d->xquat[i], xquat, 4*sizeof(float)); + mju_quat2Mat(d->xmat[i], xquat); + } + for (int i = 1; i < m->nbody; i++) { + mj_local2Global(d, d->xipos[i], d->ximat[i], m->body_ipos[i], m->body_iquat[i], i); + } + for (int g = 0; g < m->ngeom; g++) { + mj_local2Global(d, d->geom_xpos[g], d->geom_xmat[g], m->geom_pos[g], m->geom_quat[g], + m->geom_bodyid[g]); + } + for (int s = 0; s < m->nsite; s++) { + mj_local2Global(d, d->site_xpos[s], d->site_xmat[s], m->site_pos[s], m->site_quat[s], + m->site_bodyid[s]); + } +} + +// Subtree centers of mass, spatial inertias and dof motion axes in the frame +// centered at the kinematic tree's subtree COM (world orientation) +MJ_HD void mj_comPos(const MjModel* m, MjData* d) { + for (int i = 0; i < m->nbody; i++) { + for (int k = 0; k < 3; k++) { + d->subtree_com[i][k] = d->xipos[i][k]*m->body_mass[i]; + } + } + for (int i = m->nbody - 1; i > 0; i--) { + int p = m->body_parentid[i]; + for (int k = 0; k < 3; k++) { + d->subtree_com[p][k] += d->subtree_com[i][k]; + } + } + for (int i = 0; i < m->nbody; i++) { + if (m->body_subtreemass[i] < MJ_MINVAL) { + memcpy(d->subtree_com[i], d->xipos[i], 3*sizeof(float)); + continue; + } + for (int k = 0; k < 3; k++) { + d->subtree_com[i][k] /= m->body_subtreemass[i]; + } + } + memset(d->cinert[0], 0, 10*sizeof(float)); + for (int i = 1; i < m->nbody; i++) { + float* com = d->subtree_com[m->body_rootid[i]]; + float offset[3] = {d->xipos[i][0] - com[0], d->xipos[i][1] - com[1], d->xipos[i][2] - com[2]}; + mju_inertCom(d->cinert[i], m->body_inertia[i], d->ximat[i], offset, m->body_mass[i]); + } + for (int i = 1; i < m->nbody; i++) { + float* com = d->subtree_com[m->body_rootid[i]]; + for (int j = m->body_jntadr[i]; j < m->body_jntadr[i] + m->body_jntnum[i]; j++) { + int da = m->jnt_dofadr[j]; + int jtype = m->jnt_type[j]; + float offset[3] = {com[0] - d->xanchor[j][0], com[1] - d->xanchor[j][1], + com[2] - d->xanchor[j][2]}; + if (jtype == MJ_JNT_SLIDE) { + memset(d->cdof[da], 0, 6*sizeof(float)); + memcpy(d->cdof[da] + 3, d->xaxis[j], 3*sizeof(float)); + } else if (jtype == MJ_JNT_HINGE) { + memcpy(d->cdof[da], d->xaxis[j], 3*sizeof(float)); + mju_cross(d->cdof[da] + 3, d->xaxis[j], offset); + } else { + int skip = 0; + if (jtype == MJ_JNT_FREE) { + memset(d->cdof[da], 0, 18*sizeof(float)); + d->cdof[da][3] = 1.0f; + d->cdof[da + 1][4] = 1.0f; + d->cdof[da + 2][5] = 1.0f; + skip = 3; + } + for (int k = 0; k < 3; k++) { + float axis[3] = {d->xmat[i][k], d->xmat[i][k + 3], d->xmat[i][k + 6]}; + memcpy(d->cdof[da + skip + k], axis, 3*sizeof(float)); + mju_cross(d->cdof[da + skip + k] + 3, axis, offset); + } + } + } + } +} + +// Composite rigid body: dense mass matrix and its Cholesky factor +MJ_HD void mj_crb(const MjModel* m, MjData* d) { + float crb[MJ_MAX_NBODY][10]; + memcpy(crb, d->cinert, sizeof(crb)); + for (int i = m->nbody - 1; i > 0; i--) { + int p = m->body_parentid[i]; + if (p > 0) { + for (int k = 0; k < 10; k++) { + crb[p][k] += crb[i][k]; + } + } + } + memset(d->qM, 0, sizeof(d->qM)); + for (int i = 0; i < m->nv; i++) { + float buf[6]; + mju_mulInertVec(buf, crb[m->dof_bodyid[i]], d->cdof[i]); + d->qM[i][i] = m->dof_armature[i]; + for (int j = i; j >= 0; j = m->dof_parentid[j]) { + d->qM[i][j] += mju_dot6(d->cdof[j], buf); + d->qM[j][i] = d->qM[i][j]; + } + } + memcpy(d->qLD, d->qM, sizeof(d->qM)); + mju_cholFactor(&d->qLD[0][0], m->nv, MJ_MAX_NV, 1); +} + +// Body spatial velocities and dof axis derivatives (cvel x cdof) +MJ_HD void mj_comVel(const MjModel* m, MjData* d) { + memset(d->cvel[0], 0, 6*sizeof(float)); + for (int i = 1; i < m->nbody; i++) { + float cvel[6]; + memcpy(cvel, d->cvel[m->body_parentid[i]], sizeof(cvel)); + int bda = m->body_dofadr[i]; + int dofnum = m->body_dofnum[i]; + for (int j = 0; j < dofnum; j++) { + int jtype = m->jnt_type[m->dof_jntid[bda + j]]; + int n = 1; + if (jtype == MJ_JNT_FREE) { + memset(d->cdof_dot[bda], 0, 18*sizeof(float)); + for (int k = 0; k < 3; k++) { + for (int c = 0; c < 6; c++) { + cvel[c] += d->cdof[bda + k][c]*d->qvel[bda + k]; + } + } + j += 3; + n = 3; + } else if (jtype == MJ_JNT_BALL) { + n = 3; + } + for (int k = 0; k < n; k++) { + mju_crossMotion(d->cdof_dot[bda + j + k], cvel, d->cdof[bda + j + k]); + } + for (int k = 0; k < n; k++) { + for (int c = 0; c < 6; c++) { + cvel[c] += d->cdof[bda + j + k][c]*d->qvel[bda + j + k]; + } + } + j += n - 1; + } + memcpy(d->cvel[i], cvel, sizeof(cvel)); + } +} + +// Recursive Newton-Euler without acceleration +MJ_HD void mj_rne(const MjModel* m, MjData* d) { + float cacc[MJ_MAX_NBODY][6], cfrc[MJ_MAX_NBODY][6]; + memset(cacc[0], 0, 6*sizeof(float)); + cacc[0][3] = -m->opt_gravity[0]; + cacc[0][4] = -m->opt_gravity[1]; + cacc[0][5] = -m->opt_gravity[2]; + for (int i = 1; i < m->nbody; i++) { + int bda = m->body_dofadr[i]; + float tmp[6], tmp1[6]; + memcpy(cacc[i], cacc[m->body_parentid[i]], 6*sizeof(float)); + for (int j = bda; j < bda + m->body_dofnum[i]; j++) { + for (int c = 0; c < 6; c++) { + cacc[i][c] += d->cdof_dot[j][c]*d->qvel[j]; + } + } + mju_mulInertVec(cfrc[i], d->cinert[i], cacc[i]); + mju_mulInertVec(tmp, d->cinert[i], d->cvel[i]); + mju_crossForce(tmp1, d->cvel[i], tmp); + for (int c = 0; c < 6; c++) { + cfrc[i][c] += tmp1[c]; + } + } + for (int i = m->nbody - 1; i > 0; i--) { + int p = m->body_parentid[i]; + if (p > 0) { + for (int c = 0; c < 6; c++) { + cfrc[p][c] += cfrc[i][c]; + } + } + } + for (int i = 0; i < m->nv; i++) { + d->qfrc_bias[i] = mju_dot6(d->cdof[i], cfrc[m->dof_bodyid[i]]); + } +} + +// Joint springs and dof dampers +MJ_HD void mj_passive(const MjModel* m, MjData* d) { + memset(d->qfrc_passive, 0, sizeof(d->qfrc_passive)); + for (int j = 0; j < m->njnt; j++) { + float k = m->jnt_stiffness[j]; + if (k == 0.0f) { + continue; + } + int padr = m->jnt_qposadr[j]; + int dadr = m->jnt_dofadr[j]; + assert(m->jnt_type[j] >= MJ_JNT_SLIDE && "free/ball joint springs not supported"); + d->qfrc_passive[dadr] -= k*(d->qpos[padr] - m->qpos_spring[padr]); + } + for (int i = 0; i < m->nv; i++) { + d->qfrc_passive[i] -= m->dof_damping[i]*d->qvel[i]; + } +} + +// Motor actuators on hinge/slide joints: force = gear * clamped ctrl +MJ_HD void mj_actuation(const MjModel* m, MjData* d) { + memset(d->qfrc_actuator, 0, sizeof(d->qfrc_actuator)); + for (int u = 0; u < m->nu; u++) { + float ctrl = d->ctrl[u]; + if (m->actuator_ctrllimited[u]) { + ctrl = fminf(fmaxf(ctrl, m->actuator_ctrlrange[u][0]), m->actuator_ctrlrange[u][1]); + } + d->qfrc_actuator[m->jnt_dofadr[m->actuator_trnid[u]]] += m->actuator_gear[u]*ctrl; + } +} + +// Collision detection: geom pair filtering and primitive colliders + +typedef struct { + float dist; + float pos[3]; + float normal[3]; + float tangent[3]; +} MjPreContact; + +MJ_HD int mjraw_PlaneSphere(MjPreContact* con, float margin, const float* pos1, + const float* mat1, const float* pos2, float radius) { + con->normal[0] = mat1[2]; + con->normal[1] = mat1[5]; + con->normal[2] = mat1[8]; + float tmp[3] = {pos2[0] - pos1[0], pos2[1] - pos1[1], pos2[2] - pos1[2]}; + float cdist = mju_dot3(tmp, con->normal); + if (cdist > margin + radius) { + return 0; + } + con->dist = cdist - radius; + float s = -0.5f*con->dist - radius; + for (int k = 0; k < 3; k++) { + con->pos[k] = pos2[k] + s*con->normal[k]; + con->tangent[k] = 0.0f; + } + return 1; +} + +MJ_HD int mjraw_SphereSphere(MjPreContact* con, float margin, const float* pos1, + const float* mat1, float rad1, const float* pos2, const float* mat2, float rad2) { + float dif[3] = {pos1[0] - pos2[0], pos1[1] - pos2[1], pos1[2] - pos2[2]}; + float cdist_sqr = mju_dot3(dif, dif); + float min_dist = margin + rad1 + rad2; + if (cdist_sqr > min_dist*min_dist) { + return 0; + } + con->dist = sqrtf(cdist_sqr) - rad1 - rad2; + for (int k = 0; k < 3; k++) { + con->normal[k] = pos2[k] - pos1[k]; + } + if (mju_normalize3(con->normal) < MJ_MINVAL) { + float axis1[3] = {mat1[2], mat1[5], mat1[8]}; + float axis2[3] = {mat2[2], mat2[5], mat2[8]}; + mju_cross(con->normal, axis1, axis2); + mju_normalize3(con->normal); + } + float s = rad1 + 0.5f*con->dist; + for (int k = 0; k < 3; k++) { + con->pos[k] = pos1[k] + s*con->normal[k]; + con->tangent[k] = 0.0f; + } + return 1; +} + +MJ_HD int mjc_CapsuleCapsule(MjPreContact* con, float margin, const float* pos1, + const float* mat1, const float* size1, const float* pos2, const float* mat2, + const float* size2) { + float axis1[3] = {mat1[2]*size1[1], mat1[5]*size1[1], mat1[8]*size1[1]}; + float axis2[3] = {mat2[2]*size2[1], mat2[5]*size2[1], mat2[8]*size2[1]}; + float dif[3] = {pos1[0] - pos2[0], pos1[1] - pos2[1], pos1[2] - pos2[2]}; + float ma = mju_dot3(axis1, axis1); + float mb = -mju_dot3(axis1, axis2); + float mc = mju_dot3(axis2, axis2); + float u = -mju_dot3(axis1, dif); + float v = mju_dot3(axis2, dif); + float det = ma*mc - mb*mb; + float vec1[3], vec2[3]; + if (fabsf(det) >= MJ_MINVAL) { + float x1 = (mc*u - mb*v) / det; + float x2 = (ma*v - mb*u) / det; + if (x1 > 1.0f) { + x1 = 1.0f; + x2 = (v - mb) / mc; + } else if (x1 < -1.0f) { + x1 = -1.0f; + x2 = (v + mb) / mc; + } + if (x2 > 1.0f) { + x2 = 1.0f; + x1 = fminf(fmaxf((u - mb) / ma, -1.0f), 1.0f); + } else if (x2 < -1.0f) { + x2 = -1.0f; + x1 = fminf(fmaxf((u + mb) / ma, -1.0f), 1.0f); + } + for (int k = 0; k < 3; k++) { + vec1[k] = pos1[k] + x1*axis1[k]; + vec2[k] = pos2[k] + x2*axis2[k]; + } + return mjraw_SphereSphere(con, margin, vec1, mat1, size1[0], vec2, mat2, size2[0]); + } + // parallel axes: test both ends of each capsule, stop at two contacts + int n = 0; + float ends[4][2] = {{1.0f, 0.0f}, {-1.0f, 0.0f}, {0.0f, 1.0f}, {0.0f, -1.0f}}; + for (int e = 0; e < 4 && n < 2; e++) { + float x1 = ends[e][0]; + float x2 = ends[e][1]; + if (e < 2) { + x2 = fminf(fmaxf((v - x1*mb) / mc, -1.0f), 1.0f); + } else { + x1 = fminf(fmaxf((u - x2*mb) / ma, -1.0f), 1.0f); + } + for (int k = 0; k < 3; k++) { + vec1[k] = pos1[k] + x1*axis1[k]; + vec2[k] = pos2[k] + x2*axis2[k]; + } + n += mjraw_SphereSphere(con + n, margin, vec1, mat1, size1[0], vec2, mat2, size2[0]); + } + return n; +} + +MJ_HD void mj_collision(const MjModel* m, MjData* d) { + d->ncon = 0; + for (int ga = 0; ga < m->ngeom; ga++) { + for (int gb = ga + 1; gb < m->ngeom; gb++) { + int g1 = ga; + int g2 = gb; + if (m->geom_type[g1] > m->geom_type[g2]) { + g1 = gb; + g2 = ga; + } + int w1 = m->body_weldid[m->geom_bodyid[g1]]; + int w2 = m->body_weldid[m->geom_bodyid[g2]]; + if (w1 == w2 || (m->body_dofnum[w1] == 0 && m->body_dofnum[w2] == 0)) { + continue; + } + if (w1 && w2 && (w1 == m->body_weldid[m->body_parentid[w2]] + || w2 == m->body_weldid[m->body_parentid[w1]])) { + continue; + } + if (!(m->geom_contype[g1] & m->geom_conaffinity[g2]) + && !(m->geom_contype[g2] & m->geom_conaffinity[g1])) { + continue; + } + float margin = m->geom_margin[g1] + m->geom_margin[g2]; + float bound = margin + m->geom_gap[g1] + m->geom_gap[g2]; + int t1 = m->geom_type[g1]; + int t2 = m->geom_type[g2]; + if (t1 == MJ_GEOM_PLANE && t2 == MJ_GEOM_PLANE) { + continue; + } + float* pos1 = d->geom_xpos[g1]; + float* pos2 = d->geom_xpos[g2]; + float* mat1 = d->geom_xmat[g1]; + float* mat2 = d->geom_xmat[g2]; + const float* size1 = m->geom_size[g1]; + const float* size2 = m->geom_size[g2]; + float r1 = m->geom_rbound[g1]; + float r2 = m->geom_rbound[g2]; + if (t1 == MJ_GEOM_PLANE) { + float n[3] = {mat1[2], mat1[5], mat1[8]}; + float dif[3] = {pos2[0] - pos1[0], pos2[1] - pos1[1], pos2[2] - pos1[2]}; + if (mju_dot3(n, dif) > bound + r2) { + continue; + } + } else { + float dif[3] = {pos1[0] - pos2[0], pos1[1] - pos2[1], pos1[2] - pos2[2]}; + if (mju_dot3(dif, dif) > (r1 + r2 + bound)*(r1 + r2 + bound)) { + continue; + } + } + MjPreContact pre[4]; + int n = 0; + float axis2[3] = {mat2[2], mat2[5], mat2[8]}; + if (t1 == MJ_GEOM_PLANE && t2 == MJ_GEOM_SPHERE) { + n = mjraw_PlaneSphere(pre, margin, pos1, mat1, pos2, size2[0]); + } else if (t1 == MJ_GEOM_PLANE && t2 == MJ_GEOM_CAPSULE) { + // one sphere test per capsule end, frames aligned with the axis + for (int end = 1; end >= -1; end -= 2) { + float p[3] = {pos2[0] + end*size2[1]*axis2[0], pos2[1] + end*size2[1]*axis2[1], + pos2[2] + end*size2[1]*axis2[2]}; + if (mjraw_PlaneSphere(pre + n, margin, pos1, mat1, p, size2[0])) { + memcpy(pre[n].tangent, axis2, sizeof(axis2)); + n++; + } + } + } else if (t1 == MJ_GEOM_SPHERE && t2 == MJ_GEOM_SPHERE) { + n = mjraw_SphereSphere(pre, margin, pos1, mat1, size1[0], pos2, mat2, size2[0]); + } else if (t1 == MJ_GEOM_SPHERE && t2 == MJ_GEOM_CAPSULE) { + // sphere against the nearest point of the capsule segment + float vec[3] = {pos1[0] - pos2[0], pos1[1] - pos2[1], pos1[2] - pos2[2]}; + float x = fminf(fmaxf(mju_dot3(axis2, vec), -size2[1]), size2[1]); + for (int k = 0; k < 3; k++) { + vec[k] = pos2[k] + x*axis2[k]; + } + n = mjraw_SphereSphere(pre, margin, pos1, mat1, size1[0], vec, mat2, size2[0]); + } else if (t1 == MJ_GEOM_CAPSULE && t2 == MJ_GEOM_CAPSULE) { + n = mjc_CapsuleCapsule(pre, margin, pos1, mat1, size1, pos2, mat2, size2); + } else { + assert(0 && "unsupported geom pair"); + } + for (int c = 0; c < n && d->ncon < MJ_MAXCON; c++) { + MjContact* con = &d->contact[d->ncon]; + float includemargin = margin - m->geom_gap[g1] - m->geom_gap[g2]; + if (pre[c].dist >= includemargin) { + continue; + } + con->geom1 = g1; + con->geom2 = g2; + con->dist = pre[c].dist; + con->includemargin = includemargin; + memcpy(con->pos, pre[c].pos, sizeof(con->pos)); + // frame rows: normal, tangent1 (hint or a default orthogonal), tangent2 + float* frame = con->frame; + memcpy(frame, pre[c].normal, 3*sizeof(float)); + memcpy(frame + 3, pre[c].tangent, 3*sizeof(float)); + mju_normalize3(frame); + if (mju_dot3(frame + 3, frame + 3) < 0.25f) { + memset(frame + 3, 0, 3*sizeof(float)); + frame[frame[1] < 0.5f && frame[1] > -0.5f ? 4 : 5] = 1.0f; + } + float dot = mju_dot3(frame, frame + 3); + for (int k = 0; k < 3; k++) { + frame[3 + k] -= dot*frame[k]; + } + mju_normalize3(frame + 3); + mju_cross(frame + 6, frame, frame + 3); + // mix geom parameters: priority, else solmix blend, max friction/condim + int p1 = m->geom_priority[g1]; + int p2 = m->geom_priority[g2]; + float fri[3]; + if (p1 != p2) { + int g = p1 > p2 ? g1 : g2; + con->dim = m->geom_condim[g]; + memcpy(con->solref, m->geom_solref[g], sizeof(con->solref)); + memcpy(con->solimp, m->geom_solimp[g], sizeof(con->solimp)); + memcpy(fri, m->geom_friction[g], sizeof(fri)); + } else { + int c1 = m->geom_condim[g1]; + int c2 = m->geom_condim[g2]; + con->dim = c1 > c2 ? c1 : c2; + float s1 = m->geom_solmix[g1]; + float s2 = m->geom_solmix[g2]; + float mix = 0.5f; + if (s1 >= MJ_MINVAL && s2 >= MJ_MINVAL) { + mix = s1 / (s1 + s2); + } else if (s1 >= MJ_MINVAL || s2 >= MJ_MINVAL) { + mix = s1 >= MJ_MINVAL ? 1.0f : 0.0f; + } + for (int k = 0; k < 2; k++) { + con->solref[k] = mix*m->geom_solref[g1][k] + + (1.0f - mix)*m->geom_solref[g2][k]; + } + for (int k = 0; k < 5; k++) { + con->solimp[k] = mix*m->geom_solimp[g1][k] + + (1.0f - mix)*m->geom_solimp[g2][k]; + } + for (int k = 0; k < 3; k++) { + fri[k] = fmaxf(m->geom_friction[g1][k], m->geom_friction[g2][k]); + } + } + con->friction[0] = fri[0]; + con->friction[1] = fri[0]; + con->friction[2] = fri[1]; + con->friction[3] = fri[2]; + con->friction[4] = fri[2]; + d->ncon++; + } + } + } +} + +// Constraints: rows of efc_J with pos/margin, MuJoCo's impedance and +// reference acceleration, then the dual QP for the constraint forces + +// Jacobian of a world point attached to body (jacp: translation, 3 x nv rows) +MJ_HD void mj_jac(const MjModel* m, MjData* d, float jacp[3][MJ_MAX_NV], float jacr[3][MJ_MAX_NV], + const float* point, int body) { + memset(jacp, 0, 3*MJ_MAX_NV*sizeof(float)); + if (jacr) { + memset(jacr, 0, 3*MJ_MAX_NV*sizeof(float)); + } + float* com = d->subtree_com[m->body_rootid[body]]; + float offset[3] = {point[0] - com[0], point[1] - com[1], point[2] - com[2]}; + body = m->body_weldid[body]; + if (m->body_dofnum[body] == 0) { + return; + } + for (int i = m->body_dofadr[body] + m->body_dofnum[body] - 1; i >= 0; + i = m->dof_parentid[i]) { + float* cdof = d->cdof[i]; + float tmp[3]; + mju_cross(tmp, cdof, offset); + jacp[0][i] = cdof[3] + tmp[0]; + jacp[1][i] = cdof[4] + tmp[1]; + jacp[2][i] = cdof[5] + tmp[2]; + if (jacr) { + jacr[0][i] = cdof[0]; + jacr[1][i] = cdof[1]; + jacr[2][i] = cdof[2]; + } + } +} + +// Regularization R = (1 - imp)/imp * diagA and reference acceleration +// aref = -b vel - k imp (pos - margin) for `size` rows starting at row i. +MJ_HD void mj_rowParams(const MjModel* m, MjData* d, int i, int size, const float* solref, + const float* solimp, const float* diagA) { + // solimp impedance (dmin, dmax, width, midpoint, power) of the violation + float d0 = fminf(MJ_MAXIMP, fmaxf(MJ_MINIMP, solimp[0])); + float dmax = fminf(MJ_MAXIMP, fmaxf(MJ_MINIMP, solimp[1])); + float width = fmaxf(0.0f, solimp[2]); + float mid = fminf(MJ_MAXIMP, fmaxf(MJ_MINIMP, solimp[3])); + float power = fmaxf(1.0f, solimp[4]); + float x = fabsf((d->efc_pos[i] - d->efc_margin[i]) / width); + float y = x; + if (power != 1.0f && x <= mid) { + y = powf(x, power) / powf(mid, power - 1.0f); + } else if (power != 1.0f && x < 1.0f) { + y = 1.0f - powf(1.0f - x, power) / powf(1.0f - mid, power - 1.0f); + } + float imp = d0 + fminf(y, 1.0f)*(dmax - d0); + if (d0 == dmax || width <= MJ_MINVAL) { + imp = 0.5f*(d0 + dmax); + } + float ref0 = fmaxf(solref[0], 2.0f*m->opt_timestep); + float k = 1.0f / fmaxf(MJ_MINVAL, dmax*dmax*ref0*ref0*solref[1]*solref[1]); + float b = 2.0f / fmaxf(MJ_MINVAL, dmax*ref0); + for (int j = i; j < i + size; j++) { + d->efc_R[j] = fmaxf(MJ_MINVAL, (1.0f - imp)*diagA[j - i]/imp); + float vel = 0.0f; + for (int v = 0; v < m->nv; v++) { + vel += d->efc_J[j][v]*d->qvel[v]; + } + d->efc_aref[j] = -b*vel - k*imp*(d->efc_pos[j] - d->efc_margin[j]); + } +} + +MJ_HD void mj_makeConstraint(const MjModel* m, MjData* d) { + d->nefc = 0; + for (int j = 0; j < m->njnt; j++) { + if (!m->jnt_limited[j]) { + continue; + } + assert(m->jnt_type[j] == MJ_JNT_HINGE || m->jnt_type[j] == MJ_JNT_SLIDE); + float value = d->qpos[m->jnt_qposadr[j]]; + for (int side = -1; side <= 1; side += 2) { + float dist = side*(m->jnt_range[j][(side + 1)/2] - value); + if (dist >= m->jnt_margin[j]) { + continue; + } + int i = d->nefc++; + int dof = m->jnt_dofadr[j]; + memset(d->efc_J[i], 0, MJ_MAX_NV*sizeof(float)); + d->efc_J[i][dof] = -side; + d->efc_pos[i] = dist; + d->efc_margin[i] = m->jnt_margin[j]; + float diagA = m->dof_invweight0[dof]; + mj_rowParams(m, d, i, 1, m->jnt_solref[j], m->jnt_solimp[j], &diagA); + } + } + for (int c = 0; c < d->ncon; c++) { + MjContact* con = &d->contact[c]; + int b1 = m->geom_bodyid[con->geom1]; + int b2 = m->geom_bodyid[con->geom2]; + int dim = con->dim; + int rows = dim == 1 ? 1 : 2*(dim - 1); + con->efc_address = -1; + if (d->nefc + rows > MJ_MAXEFC) { + break; + } + // Jacobian difference (body2 - body1) at the contact point, rotated + // into the contact frame: rows normal, tangent1, tangent2, and for + // condim > 3 the rotational rows torsion, roll1, roll2 + float jac1[3][MJ_MAX_NV], jac2[3][MJ_MAX_NV], jac[6][MJ_MAX_NV]; + float jacr1[3][MJ_MAX_NV], jacr2[3][MJ_MAX_NV]; + mj_jac(m, d, jac1, dim > 3 ? jacr1 : NULL, con->pos, b1); + mj_jac(m, d, jac2, dim > 3 ? jacr2 : NULL, con->pos, b2); + for (int r = 0; r < 3; r++) { + for (int v = 0; v < m->nv; v++) { + float dp = jac2[0][v] - jac1[0][v]; + float dq = jac2[1][v] - jac1[1][v]; + float dr = jac2[2][v] - jac1[2][v]; + jac[r][v] = con->frame[3*r]*dp + con->frame[3*r + 1]*dq + con->frame[3*r + 2]*dr; + if (dim > 3) { + dp = jacr2[0][v] - jacr1[0][v]; + dq = jacr2[1][v] - jacr1[1][v]; + dr = jacr2[2][v] - jacr1[2][v]; + jac[3 + r][v] = con->frame[3*r]*dp + con->frame[3*r + 1]*dq + + con->frame[3*r + 2]*dr; + } + } + } + int i0 = d->nefc; + con->efc_address = i0; + float tran = m->body_invweight0[b1][0] + m->body_invweight0[b2][0]; + float rot = m->body_invweight0[b1][1] + m->body_invweight0[b2][1]; + float diagA[10]; + if (dim == 1) { + memcpy(d->efc_J[i0], jac[0], MJ_MAX_NV*sizeof(float)); + diagA[0] = tran; + } else { + for (int k = 1; k < dim; k++) { + float fri = con->friction[k - 1]; + for (int v = 0; v < m->nv; v++) { + d->efc_J[i0 + 2*(k - 1)][v] = jac[0][v] + fri*jac[k][v]; + d->efc_J[i0 + 2*(k - 1) + 1][v] = jac[0][v] - fri*jac[k][v]; + } + diagA[2*(k - 1)] = tran + fri*fri*(k < 3 ? tran : rot); + diagA[2*(k - 1) + 1] = diagA[2*(k - 1)]; + } + } + for (int i = i0; i < i0 + rows; i++) { + d->efc_pos[i] = con->dist; + d->efc_margin[i] = con->includemargin; + } + d->nefc += rows; + mj_rowParams(m, d, i0, rows, con->solref, con->solimp, diagA); + if (dim > 1) { + // pyramidal cone: common R matching the friction impedance of the + // elliptic model, R1 = R0 / impratio + float r1 = d->efc_R[i0] / fmaxf(MJ_MINVAL, m->opt_impratio); + float mu = con->friction[0]*sqrtf(r1 / d->efc_R[i0]); + float rpy = 2.0f*mu*mu*d->efc_R[i0]; + for (int i = i0; i < i0 + rows; i++) { + d->efc_R[i] = rpy; + } + } + } +} + +// Constraint forces: min 1/2 f^T (A + R) f - f^T (aref - J qacc_smooth) with +// f >= 0 and A = J M^-1 J^T, by active-set (Lawson-Hanson NNLS) on the free +// set. Sets efc_force, qfrc_constraint and qacc. The dense matrices are the +// scratch bound by mj_makeData: efc_AR = A + R, efc_ARfree its factored +// free-set block and efc_MinvJT = M^-1 J^T, indexed with stride s. +MJ_HD void mj_solveConstraint(const MjModel* m, MjData* d) { + int nefc = d->nefc; + int s = d->efc_stride; + memcpy(d->qacc, d->qacc_smooth, sizeof(d->qacc)); + memset(d->qfrc_constraint, 0, sizeof(d->qfrc_constraint)); + if (nefc == 0) { + return; + } + float* MinvJT = d->efc_MinvJT; + float* G = d->efc_AR; + float* Gp = d->efc_ARfree; + float b[MJ_MAXEFC], z[MJ_MAXEFC], dqacc[MJ_MAX_NV] = {0}; + float* fc = d->efc_force; + int isfree[MJ_MAXEFC], idx[MJ_MAXEFC]; + for (int i = 0; i < nefc; i++) { + float row[MJ_MAX_NV]; + memcpy(row, d->efc_J[i], sizeof(row)); + mju_cholSolve(&d->qLD[0][0], m->nv, MJ_MAX_NV, 1, row); + b[i] = d->efc_aref[i]; + for (int v = 0; v < m->nv; v++) { + MinvJT[(i*MJ_MAX_NV + v)*s] = row[v]; + b[i] -= d->efc_J[i][v]*d->qacc_smooth[v]; + } + for (int j = 0; j <= i; j++) { + float g = 0.0f; + for (int v = 0; v < m->nv; v++) { + g += d->efc_J[i][v]*MinvJT[(j*MJ_MAX_NV + v)*s]; + } + G[(i*MJ_MAXEFC + j)*s] = g; + G[(j*MJ_MAXEFC + i)*s] = g; + } + G[(i*MJ_MAXEFC + i)*s] += d->efc_R[i]; + fc[i] = 0.0f; + isfree[i] = 0; + } + for (int it = 0; it < 3*nefc + 10; it++) { + // most violated bound row (negative gradient) joins the free set + int best = -1; + float gbest = 0.0f; + for (int i = 0; i < nefc; i++) { + if (isfree[i]) { + continue; + } + float g = -b[i]; + for (int v = 0; v < m->nv; v++) { + g += d->efc_J[i][v]*dqacc[v]; + } + if (g < gbest - 1e-3f*(1.0f + 0.01f*fabsf(b[i]))) { + gbest = g; + best = i; + } + } + if (best < 0) { + break; + } + isfree[best] = 1; + // solve on the free set; step until a free row hits zero and drop it + for (;;) { + int np = 0; + for (int i = 0; i < nefc; i++) { + if (isfree[i]) { + idx[np++] = i; + } + } + for (int a = 0; a < np; a++) { + z[a] = b[idx[a]]; + for (int c = 0; c < np; c++) { + Gp[(a*MJ_MAXEFC + c)*s] = G[(idx[a]*MJ_MAXEFC + idx[c])*s]; + } + } + mju_cholFactor(Gp, np, MJ_MAXEFC, s); + mju_cholSolve(Gp, np, MJ_MAXEFC, s, z); + float alpha = 1.0f; + for (int a = 0; a < np; a++) { + float fi = fc[idx[a]]; + if (z[a] <= 0.0f && fi > z[a]) { + alpha = fminf(alpha, fi / (fi - z[a])); + } + } + for (int a = 0; a < np; a++) { + int i = idx[a]; + float fi = fc[i]; + fc[i] = fi + alpha*(z[a] - fi); + if (z[a] <= 0.0f && fi <= alpha*(fi - z[a])*(1.0f + 1e-5f)) { + fc[i] = 0.0f; + isfree[i] = 0; + } + } + if (alpha >= 1.0f) { + break; + } + } + memset(dqacc, 0, sizeof(dqacc)); + for (int i = 0; i < nefc; i++) { + for (int v = 0; v < m->nv; v++) { + dqacc[v] += fc[i]*MinvJT[(i*MJ_MAX_NV + v)*s]; + } + } + } + for (int i = 0; i < nefc; i++) { + for (int v = 0; v < m->nv; v++) { + d->qfrc_constraint[v] += fc[i]*d->efc_J[i][v]; + d->qacc[v] += fc[i]*MinvJT[(i*MJ_MAX_NV + v)*s]; + } + } +} + +// Body accelerations and interaction forces including constraint forces: +// cacc, cfrc_int and cfrc_ext (torque:force in the subtree COM frame; contact +// forces decoded from the pyramid rows) +MJ_HD void mj_rnePostConstraint(const MjModel* m, MjData* d) { + memset(d->cfrc_ext, 0, sizeof(d->cfrc_ext)); + for (int c = 0; c < d->ncon; c++) { + MjContact* con = &d->contact[c]; + int adr = con->efc_address; + if (adr < 0) { + continue; + } + float lfrc[6] = {0}; + if (con->dim == 1) { + lfrc[0] = d->efc_force[adr]; + } else { + for (int k = 0; k < con->dim - 1; k++) { + lfrc[0] += d->efc_force[adr + 2*k] + d->efc_force[adr + 2*k + 1]; + lfrc[k + 1] = (d->efc_force[adr + 2*k] - d->efc_force[adr + 2*k + 1]) + *con->friction[k]; + } + } + float cfrc[6]; + mju_mulMatTVec3(cfrc, con->frame, lfrc + 3); + mju_mulMatTVec3(cfrc + 3, con->frame, lfrc); + int bodies[2] = {m->geom_bodyid[con->geom1], m->geom_bodyid[con->geom2]}; + for (int side = 0; side < 2; side++) { + int k = bodies[side]; + if (k == 0) { + continue; + } + float* com = d->subtree_com[m->body_rootid[k]]; + float dif[3] = {com[0] - con->pos[0], com[1] - con->pos[1], com[2] - con->pos[2]}; + float cros[3]; + mju_cross(cros, dif, cfrc + 3); + float sign = side ? 1.0f : -1.0f; + for (int i = 0; i < 3; i++) { + d->cfrc_ext[k][i] += sign*(cfrc[i] - cros[i]); + d->cfrc_ext[k][i + 3] += sign*cfrc[i + 3]; + } + } + } + memset(d->cacc[0], 0, 6*sizeof(float)); + memset(d->cfrc_int[0], 0, 6*sizeof(float)); + d->cacc[0][3] = -m->opt_gravity[0]; + d->cacc[0][4] = -m->opt_gravity[1]; + d->cacc[0][5] = -m->opt_gravity[2]; + for (int j = 1; j < m->nbody; j++) { + int bda = m->body_dofadr[j]; + float cfrc_body[6], cfrc_corr[6], cfrc[6]; + memcpy(d->cacc[j], d->cacc[m->body_parentid[j]], 6*sizeof(float)); + for (int i = bda; i < bda + m->body_dofnum[j]; i++) { + for (int c = 0; c < 6; c++) { + d->cacc[j][c] += d->cdof_dot[i][c]*d->qvel[i] + d->cdof[i][c]*d->qacc[i]; + } + } + mju_mulInertVec(cfrc_body, d->cinert[j], d->cacc[j]); + mju_mulInertVec(cfrc_corr, d->cinert[j], d->cvel[j]); + mju_crossForce(cfrc, d->cvel[j], cfrc_corr); + for (int c = 0; c < 6; c++) { + d->cfrc_int[j][c] = cfrc_body[c] + cfrc[c] - d->cfrc_ext[j][c]; + } + } + for (int j = m->nbody - 1; j > 0; j--) { + for (int c = 0; c < 6; c++) { + d->cfrc_int[m->body_parentid[j]][c] += d->cfrc_int[j][c]; + } + } +} + +// Forward dynamics: qacc and all intermediate quantities from qpos, qvel, ctrl +MJ_HD void mj_forward(const MjModel* m, MjData* d) { + mj_kinematics(m, d); + mj_comPos(m, d); + mj_crb(m, d); + mj_collision(m, d); + mj_comVel(m, d); + mj_passive(m, d); + mj_makeConstraint(m, d); + mj_rne(m, d); + mj_actuation(m, d); + for (int i = 0; i < m->nv; i++) { + d->qfrc_smooth[i] = d->qfrc_passive[i] - d->qfrc_bias[i] + d->qfrc_actuator[i]; + } + memcpy(d->qacc_smooth, d->qfrc_smooth, sizeof(d->qacc_smooth)); + mju_cholSolve(&d->qLD[0][0], m->nv, MJ_MAX_NV, 1, d->qacc_smooth); + mj_solveConstraint(m, d); +} + +MJ_HD void mj_integratePos(const MjModel* m, float* qpos, const float* qvel, float dt) { + for (int j = 0; j < m->njnt; j++) { + int padr = m->jnt_qposadr[j]; + int vadr = m->jnt_dofadr[j]; + int jtype = m->jnt_type[j]; + if (jtype == MJ_JNT_FREE) { + for (int k = 0; k < 3; k++) { + qpos[padr + k] += dt*qvel[vadr + k]; + } + mju_quatIntegrate(qpos + padr + 3, qvel + vadr + 3, dt); + } else if (jtype == MJ_JNT_BALL) { + mju_quatIntegrate(qpos + padr, qvel + vadr, dt); + } else { + qpos[padr] += dt*qvel[vadr]; + } + } +} + +// Semi-implicit Euler with implicit joint damping: (M + h B) qacc' = M qacc +MJ_HD void mj_Euler(const MjModel* m, MjData* d) { + float h = m->opt_timestep; + float qacc[MJ_MAX_NV]; + float MhB[MJ_MAX_NV][MJ_MAX_NV]; + memcpy(MhB, d->qM, sizeof(MhB)); + for (int i = 0; i < m->nv; i++) { + MhB[i][i] += h*m->dof_damping[i]; + qacc[i] = d->qfrc_smooth[i] + d->qfrc_constraint[i]; + } + mju_cholFactor(&MhB[0][0], m->nv, MJ_MAX_NV, 1); + mju_cholSolve(&MhB[0][0], m->nv, MJ_MAX_NV, 1, qacc); + for (int i = 0; i < m->nv; i++) { + d->qvel[i] += h*qacc[i]; + } + mj_integratePos(m, d->qpos, d->qvel, h); + d->time += h; +} + +// Explicit RK4 (mj_RungeKutta with N=4) +MJ_HD void mj_RungeKutta4(const MjModel* m, MjData* d) { + float h = m->opt_timestep; + float A[3] = {0.5f, 0.5f, 1.0f}; + float B[4] = {1.0f/6.0f, 1.0f/3.0f, 1.0f/3.0f, 1.0f/6.0f}; + float X[4][MJ_MAX_NQ + MJ_MAX_NV], F[4][MJ_MAX_NV]; + float qpos0[MJ_MAX_NQ], time0 = d->time; + memcpy(qpos0, d->qpos, sizeof(qpos0)); + memcpy(X[0], d->qpos, sizeof(qpos0)); + memcpy(X[0] + m->nq, d->qvel, sizeof(d->qvel)); + memcpy(F[0], d->qacc, sizeof(d->qacc)); + for (int i = 1; i < 4; i++) { + // stage i uses only stage i-1 with weight A[i-1] + float dv[MJ_MAX_NV], da[MJ_MAX_NV]; + for (int v = 0; v < m->nv; v++) { + dv[v] = A[i - 1]*X[i - 1][m->nq + v]; + da[v] = A[i - 1]*F[i - 1][v]; + } + memcpy(d->qpos, qpos0, sizeof(qpos0)); + mj_integratePos(m, d->qpos, dv, h); + for (int v = 0; v < m->nv; v++) { + d->qvel[v] = X[0][m->nq + v] + h*da[v]; + } + mj_forward(m, d); + memcpy(X[i], d->qpos, sizeof(qpos0)); + memcpy(X[i] + m->nq, d->qvel, sizeof(d->qvel)); + memcpy(F[i], d->qacc, sizeof(d->qacc)); + } + float dv[MJ_MAX_NV], da[MJ_MAX_NV]; + for (int v = 0; v < m->nv; v++) { + dv[v] = 0.0f; + da[v] = 0.0f; + for (int j = 0; j < 4; j++) { + dv[v] += B[j]*X[j][m->nq + v]; + da[v] += B[j]*F[j][v]; + } + } + memcpy(d->qpos, qpos0, sizeof(qpos0)); + for (int v = 0; v < m->nv; v++) { + d->qvel[v] = X[0][m->nq + v] + h*da[v]; + } + mj_integratePos(m, d->qpos, dv, h); + d->time = time0 + h; +} + +MJ_HD void mj_step(const MjModel* m, MjData* d) { + mj_forward(m, d); + if (m->opt_integrator == MJ_INT_RK4) { + mj_RungeKutta4(m, d); + } else { + assert(m->opt_integrator == MJ_INT_EULER); + mj_Euler(m, d); + } +} + +MJ_HD void mj_resetData(const MjModel* m, MjData* d) { + memset(&d->time, 0, sizeof(MjData) - offsetof(MjData, time)); + memcpy(d->qpos, m->qpos0, sizeof(d->qpos)); +} diff --git a/ocean/mujoco/render.h b/ocean/mujoco/render.h new file mode 100644 index 0000000000..b0b974228e --- /dev/null +++ b/ocean/mujoco/render.h @@ -0,0 +1,69 @@ +// Raylib 3D rendering of a physics.h model: geoms, contacts and a tracking +// camera. MuJoCo is z-up, raylib y-up: (x, y, z) -> (x, z, -y). +#include +#include "raylib.h" +#include "rlgl.h" + +const Color MJ_BACKGROUND = (Color){6, 24, 24, 255}; +const Color MJ_BODY = (Color){0, 187, 187, 255}; +const Color MJ_CONTACT = (Color){187, 0, 0, 255}; + +Vector3 mj_rl(const float* p) { + return (Vector3){p[0], p[2], -p[1]}; +} + +// Draw one frame from qpos: window setup, camera `back` behind and `up` above +// the target (MuJoCo coords), geoms, contacts and a HUD line +void mj_render(const MjModel* m, MjData* d, const char* title, const float* target, float back, + float up, const char* hud) { + if (!IsWindowReady()) { + InitWindow(1280, 720, title); + SetTargetFPS(60); + } + if (IsKeyDown(KEY_ESCAPE)) { + exit(0); + } + mj_kinematics(m, d); + mj_collision(m, d); + float eye[3] = {target[0], target[1] - back, target[2] + up}; + Camera3D camera = {.position = mj_rl(eye), .target = mj_rl(target), .up = {0, 1, 0}, + .fovy = 45.0f, .projection = CAMERA_PERSPECTIVE}; + BeginDrawing(); + ClearBackground(MJ_BACKGROUND); + BeginMode3D(camera); + for (int g = 0; g < m->ngeom; g++) { + const float* size = m->geom_size[g]; + float* pos = d->geom_xpos[g]; + float* mat = d->geom_xmat[g]; + float a[3] = {pos[0] - size[1]*mat[2], pos[1] - size[1]*mat[5], pos[2] - size[1]*mat[8]}; + float b[3] = {pos[0] + size[1]*mat[2], pos[1] + size[1]*mat[5], pos[2] + size[1]*mat[8]}; + int type = m->geom_type[g]; + if (type == MJ_GEOM_PLANE) { + rlPushMatrix(); + rlTranslatef(pos[0], pos[2], -pos[1]); + DrawGrid(400, 1.0f); + rlPopMatrix(); + } else if (type == MJ_GEOM_SPHERE) { + DrawSphere(mj_rl(pos), size[0], MJ_BODY); + } else if (type == MJ_GEOM_CAPSULE) { + DrawCapsule(mj_rl(a), mj_rl(b), size[0], 12, 4, MJ_BODY); + } else if (type == MJ_GEOM_CYLINDER) { + DrawCylinderEx(mj_rl(a), mj_rl(b), size[0], size[0], 16, MJ_BODY); + } else if (type == MJ_GEOM_BOX) { + // column-major OpenGL matrix of P R P^T with P the z-up to y-up map + float t[16] = {mat[0], mat[6], -mat[3], 0.0f, mat[2], mat[8], -mat[5], 0.0f, + -mat[1], -mat[7], mat[4], 0.0f, pos[0], pos[2], -pos[1], 1.0f}; + rlPushMatrix(); + rlMultMatrixf(t); + DrawCube((Vector3){0, 0, 0}, 2.0f*size[0], 2.0f*size[2], 2.0f*size[1], MJ_BODY); + rlPopMatrix(); + } + } + for (int c = 0; c < d->ncon; c++) { + DrawSphere(mj_rl(d->contact[c].pos), 0.02f, MJ_CONTACT); + } + EndMode3D(); + DrawText(hud, 20, 20, 20, RAYWHITE); + EndDrawing(); + puf_web_vsync(); +} diff --git a/resources/mujoco/ant.xml b/resources/mujoco/ant.xml new file mode 100644 index 0000000000..ee4d679981 --- /dev/null +++ b/resources/mujoco/ant.xml @@ -0,0 +1,81 @@ + + + diff --git a/resources/mujoco/half_cheetah.xml b/resources/mujoco/half_cheetah.xml new file mode 100644 index 0000000000..338c2e87a7 --- /dev/null +++ b/resources/mujoco/half_cheetah.xml @@ -0,0 +1,96 @@ + + + + + + + + + + diff --git a/resources/mujoco/hopper.xml b/resources/mujoco/hopper.xml new file mode 100644 index 0000000000..cc3fbc849e --- /dev/null +++ b/resources/mujoco/hopper.xml @@ -0,0 +1,53 @@ + + + + + + + + + diff --git a/resources/mujoco/humanoid.xml b/resources/mujoco/humanoid.xml new file mode 100644 index 0000000000..a263bdbcf5 --- /dev/null +++ b/resources/mujoco/humanoid.xml @@ -0,0 +1,121 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/resources/mujoco/walker2d.xml b/resources/mujoco/walker2d.xml new file mode 100644 index 0000000000..12baef14b7 --- /dev/null +++ b/resources/mujoco/walker2d.xml @@ -0,0 +1,68 @@ + + + + + + + + From a6c7cf9410be0870adfaf4bbbcbb6fec3f5174cb Mon Sep 17 00:00:00 2001 From: vyeoms Date: Mon, 7 Sep 2026 15:11:20 +0100 Subject: [PATCH 2/5] MuJoCo cleanup --- config/mjc_ant.ini | 93 ++++++++++++++++---- config/mjc_half_cheetah.ini | 91 ++++++++++++++++---- config/mjc_hopper.ini | 93 ++++++++++++++++---- config/mjc_humanoid.ini | 89 +++++++++++++++---- config/mjc_walker2d.ini | 93 ++++++++++++++++---- ocean/mujoco/backend.h | 81 ++++++++++------- ocean/mujoco/mjc_ant.h | 34 ++++---- ocean/mujoco/mjc_half_cheetah.h | 35 +++----- ocean/mujoco/mjc_hopper.h | 37 +++----- ocean/mujoco/mjc_humanoid.h | 51 +++++------ ocean/mujoco/mjc_walker2d.h | 39 ++++----- ocean/mujoco/mjcf2bin.py | 4 + ocean/mujoco/physics.h | 148 ++++++++++++++++---------------- ocean/mujoco/render.h | 10 ++- 14 files changed, 593 insertions(+), 305 deletions(-) diff --git a/config/mjc_ant.ini b/config/mjc_ant.ini index 87dd196491..7bcba27dbd 100644 --- a/config/mjc_ant.ini +++ b/config/mjc_ant.ini @@ -2,14 +2,12 @@ env_name = mjc_ant [vec] -total_agents = 2048 -# GPU env builds (--cu) require 1 buffer; the CPU vec is ~1.5x faster with 4 +total_agents = 256 num_buffers = 1 num_threads = 16 [env] model = resources/mujoco/ant.bin -# Gymnasium Ant-v5 defaults max_steps = 1000 reset_noise_scale = 0.1 forward_reward_weight = 1.0 @@ -18,31 +16,92 @@ contact_cost_weight = 0.0005 healthy_reward = 1.0 [policy] -hidden_size = 256 +hidden_size = 512 num_layers = 2 expansion_factor = 1 [train] gpus = 1 -seed = 42 -total_timesteps = 100000000 -learning_rate = 0.005 +total_timesteps = 10000000 +learning_rate = 0.003 anneal_lr = 1 min_lr_ratio = 0 -gamma = 0.99 -gae_lambda = 0.95 -replay_ratio = 2 -clip_coef = 0.2 -vf_coef = 0.5 +gamma = 0.978 +gae_lambda = 0.944 +replay_ratio = 5.7 +clip_coef = 0.23 +vf_coef = 1.6 vf_clip_coef = 1.0 -max_grad_norm = 0.5 -ent_coef = 0.0 -momentum = 0.9 -minibatch_size = 8192 -horizon = 64 +max_grad_norm = 1.05 +ent_coef = 0.00031 +momentum = 0.99 +minibatch_size = 2048 +horizon = 32 vtrace_rho_clip = 1.0 vtrace_c_clip = 1.0 [sweep] metric = score goal = maximize +max_runs = 30 + +[sweep.train.total_timesteps] +min = 1e7 +max = 1.2e7 + +[sweep.policy.hidden_size] +min = 256 +max = 512 + +[sweep.policy.num_layers] +distribution = int_uniform +min = 2 +max = 3 + +[sweep.vec.total_agents] +min = 128 +max = 1024 + +[sweep.train.horizon] +min = 32 +max = 256 + +[sweep.train.minibatch_size] +min = 512 +max = 4096 + +[sweep.train.replay_ratio] +min = 2 +max = 6 + +[sweep.train.learning_rate] +min = 0.0005 +max = 0.02 + +[sweep.train.momentum] +min = 0.5 +max = 0.99 + +[sweep.train.ent_coef] +min = 0.00001 +max = 0.01 + +[sweep.train.gamma] +min = 0.88 +max = 0.99 + +[sweep.train.gae_lambda] +min = 0.8 +max = 0.95 + +[sweep.train.clip_coef] +min = 0.1 +max = 0.35 + +[sweep.train.vf_coef] +min = 0.5 +max = 2.0 + +[sweep.train.vf_clip_coef] +min = 1.0 +max = 5.0 diff --git a/config/mjc_half_cheetah.ini b/config/mjc_half_cheetah.ini index e7978c813e..9b6b47ccc9 100644 --- a/config/mjc_half_cheetah.ini +++ b/config/mjc_half_cheetah.ini @@ -2,14 +2,12 @@ env_name = mjc_half_cheetah [vec] -total_agents = 2048 -# GPU env builds (--cu) require 1 buffer; the CPU vec is ~1.5x faster with 4 +total_agents = 512 num_buffers = 1 num_threads = 16 [env] model = resources/mujoco/half_cheetah.bin -# Gymnasium HalfCheetah-v5 defaults. Episode ends (terminal) at max_steps. max_steps = 1000 reset_noise_scale = 0.1 forward_reward_weight = 1.0 @@ -22,25 +20,86 @@ expansion_factor = 1 [train] gpus = 1 -seed = 42 -total_timesteps = 100000000 -learning_rate = 0.005 +total_timesteps = 10000000 +learning_rate = 0.0036 anneal_lr = 1 min_lr_ratio = 0 -gamma = 0.99 -gae_lambda = 0.95 -replay_ratio = 2 -clip_coef = 0.2 +gamma = 0.975 +gae_lambda = 0.922 +replay_ratio = 4.9 +clip_coef = 0.11 vf_coef = 0.5 -vf_clip_coef = 1.0 -max_grad_norm = 0.5 -ent_coef = 0.0 -momentum = 0.9 -minibatch_size = 8192 -horizon = 64 +vf_clip_coef = 4.1 +max_grad_norm = 1.2 +ent_coef = 0.0017 +momentum = 0.955 +minibatch_size = 1024 +horizon = 32 vtrace_rho_clip = 1.0 vtrace_c_clip = 1.0 [sweep] metric = score goal = maximize +max_runs = 30 + +[sweep.train.total_timesteps] +min = 1e7 +max = 1.2e7 + +[sweep.policy.hidden_size] +min = 256 +max = 512 + +[sweep.policy.num_layers] +distribution = int_uniform +min = 2 +max = 3 + +[sweep.vec.total_agents] +min = 128 +max = 1024 + +[sweep.train.horizon] +min = 32 +max = 256 + +[sweep.train.minibatch_size] +min = 512 +max = 4096 + +[sweep.train.replay_ratio] +min = 2 +max = 6 + +[sweep.train.learning_rate] +min = 0.0005 +max = 0.02 + +[sweep.train.momentum] +min = 0.5 +max = 0.99 + +[sweep.train.ent_coef] +min = 0.00001 +max = 0.01 + +[sweep.train.gamma] +min = 0.88 +max = 0.99 + +[sweep.train.gae_lambda] +min = 0.8 +max = 0.95 + +[sweep.train.clip_coef] +min = 0.1 +max = 0.35 + +[sweep.train.vf_coef] +min = 0.5 +max = 2.0 + +[sweep.train.vf_clip_coef] +min = 1.0 +max = 5.0 diff --git a/config/mjc_hopper.ini b/config/mjc_hopper.ini index e256ed8120..634528ca9c 100644 --- a/config/mjc_hopper.ini +++ b/config/mjc_hopper.ini @@ -2,14 +2,12 @@ env_name = mjc_hopper [vec] -total_agents = 2048 -# GPU env builds (--cu) require 1 buffer; the CPU vec is ~1.5x faster with 4 +total_agents = 256 num_buffers = 1 num_threads = 16 [env] model = resources/mujoco/hopper.bin -# Gymnasium Hopper-v5 defaults max_steps = 1000 reset_noise_scale = 0.005 forward_reward_weight = 1.0 @@ -23,25 +21,86 @@ expansion_factor = 1 [train] gpus = 1 -seed = 42 -total_timesteps = 100000000 -learning_rate = 0.005 +total_timesteps = 10000000 +learning_rate = 0.0017 anneal_lr = 1 min_lr_ratio = 0 -gamma = 0.99 -gae_lambda = 0.95 -replay_ratio = 2 -clip_coef = 0.2 -vf_coef = 0.5 -vf_clip_coef = 1.0 -max_grad_norm = 0.5 -ent_coef = 0.0 -momentum = 0.9 -minibatch_size = 8192 -horizon = 64 +gamma = 0.983 +gae_lambda = 0.939 +replay_ratio = 4.5 +clip_coef = 0.26 +vf_coef = 1.8 +vf_clip_coef = 2.6 +max_grad_norm = 0.97 +ent_coef = 0.00096 +momentum = 0.974 +minibatch_size = 512 +horizon = 128 vtrace_rho_clip = 1.0 vtrace_c_clip = 1.0 [sweep] metric = score goal = maximize +max_runs = 30 + +[sweep.train.total_timesteps] +min = 1e7 +max = 1.2e7 + +[sweep.policy.hidden_size] +min = 256 +max = 512 + +[sweep.policy.num_layers] +distribution = int_uniform +min = 2 +max = 3 + +[sweep.vec.total_agents] +min = 128 +max = 1024 + +[sweep.train.horizon] +min = 32 +max = 256 + +[sweep.train.minibatch_size] +min = 512 +max = 4096 + +[sweep.train.replay_ratio] +min = 2 +max = 6 + +[sweep.train.learning_rate] +min = 0.0005 +max = 0.02 + +[sweep.train.momentum] +min = 0.5 +max = 0.99 + +[sweep.train.ent_coef] +min = 0.00001 +max = 0.01 + +[sweep.train.gamma] +min = 0.88 +max = 0.99 + +[sweep.train.gae_lambda] +min = 0.8 +max = 0.95 + +[sweep.train.clip_coef] +min = 0.1 +max = 0.35 + +[sweep.train.vf_coef] +min = 0.5 +max = 2.0 + +[sweep.train.vf_clip_coef] +min = 1.0 +max = 5.0 diff --git a/config/mjc_humanoid.ini b/config/mjc_humanoid.ini index 8d78c85ea9..9a0a94d34b 100644 --- a/config/mjc_humanoid.ini +++ b/config/mjc_humanoid.ini @@ -2,14 +2,12 @@ env_name = mjc_humanoid [vec] -total_agents = 2048 -# GPU env builds (--cu) require 1 buffer; the CPU vec is ~1.5x faster with 4 +total_agents = 512 num_buffers = 1 num_threads = 16 [env] model = resources/mujoco/humanoid.bin -# Gymnasium Humanoid-v5 defaults (contact cost clamped at 10 in the env) max_steps = 1000 reset_noise_scale = 0.01 forward_reward_weight = 1.25 @@ -24,25 +22,86 @@ expansion_factor = 1 [train] gpus = 1 -seed = 42 -total_timesteps = 100000000 -learning_rate = 0.005 +total_timesteps = 10000000 +learning_rate = 0.00093 anneal_lr = 1 min_lr_ratio = 0 gamma = 0.99 -gae_lambda = 0.95 -replay_ratio = 2 -clip_coef = 0.2 -vf_coef = 0.5 -vf_clip_coef = 1.0 +gae_lambda = 0.8 +replay_ratio = 5.7 +clip_coef = 0.15 +vf_coef = 2.0 +vf_clip_coef = 4.4 max_grad_norm = 0.5 -ent_coef = 0.0 -momentum = 0.9 -minibatch_size = 8192 -horizon = 64 +ent_coef = 0.000022 +momentum = 0.965 +minibatch_size = 512 +horizon = 128 vtrace_rho_clip = 1.0 vtrace_c_clip = 1.0 [sweep] metric = score goal = maximize +max_runs = 25 + +[sweep.train.total_timesteps] +min = 1e7 +max = 1.2e7 + +[sweep.policy.hidden_size] +min = 256 +max = 512 + +[sweep.policy.num_layers] +distribution = int_uniform +min = 2 +max = 3 + +[sweep.vec.total_agents] +min = 128 +max = 1024 + +[sweep.train.horizon] +min = 32 +max = 256 + +[sweep.train.minibatch_size] +min = 512 +max = 4096 + +[sweep.train.replay_ratio] +min = 2 +max = 6 + +[sweep.train.learning_rate] +min = 0.0005 +max = 0.02 + +[sweep.train.momentum] +min = 0.5 +max = 0.99 + +[sweep.train.ent_coef] +min = 0.00001 +max = 0.01 + +[sweep.train.gamma] +min = 0.88 +max = 0.99 + +[sweep.train.gae_lambda] +min = 0.8 +max = 0.95 + +[sweep.train.clip_coef] +min = 0.1 +max = 0.35 + +[sweep.train.vf_coef] +min = 0.5 +max = 2.0 + +[sweep.train.vf_clip_coef] +min = 1.0 +max = 5.0 diff --git a/config/mjc_walker2d.ini b/config/mjc_walker2d.ini index fad1459c43..e72f6d626d 100644 --- a/config/mjc_walker2d.ini +++ b/config/mjc_walker2d.ini @@ -2,14 +2,12 @@ env_name = mjc_walker2d [vec] -total_agents = 2048 -# GPU env builds (--cu) require 1 buffer; the CPU vec is ~1.5x faster with 4 +total_agents = 256 num_buffers = 1 num_threads = 16 [env] model = resources/mujoco/walker2d.bin -# Gymnasium Walker2d-v5 defaults max_steps = 1000 reset_noise_scale = 0.005 forward_reward_weight = 1.0 @@ -23,25 +21,86 @@ expansion_factor = 1 [train] gpus = 1 -seed = 42 -total_timesteps = 100000000 -learning_rate = 0.005 +total_timesteps = 10000000 +learning_rate = 0.004 anneal_lr = 1 min_lr_ratio = 0 -gamma = 0.99 -gae_lambda = 0.95 -replay_ratio = 2 -clip_coef = 0.2 -vf_coef = 0.5 -vf_clip_coef = 1.0 -max_grad_norm = 0.5 -ent_coef = 0.0 -momentum = 0.9 -minibatch_size = 8192 -horizon = 64 +gamma = 0.971 +gae_lambda = 0.944 +replay_ratio = 4.7 +clip_coef = 0.31 +vf_coef = 1.7 +vf_clip_coef = 3.5 +max_grad_norm = 0.77 +ent_coef = 0.0035 +momentum = 0.932 +minibatch_size = 1024 +horizon = 128 vtrace_rho_clip = 1.0 vtrace_c_clip = 1.0 [sweep] metric = score goal = maximize +max_runs = 30 + +[sweep.train.total_timesteps] +min = 1e7 +max = 1.2e7 + +[sweep.policy.hidden_size] +min = 256 +max = 512 + +[sweep.policy.num_layers] +distribution = int_uniform +min = 2 +max = 3 + +[sweep.vec.total_agents] +min = 128 +max = 1024 + +[sweep.train.horizon] +min = 32 +max = 256 + +[sweep.train.minibatch_size] +min = 512 +max = 4096 + +[sweep.train.replay_ratio] +min = 2 +max = 6 + +[sweep.train.learning_rate] +min = 0.0005 +max = 0.02 + +[sweep.train.momentum] +min = 0.5 +max = 0.99 + +[sweep.train.ent_coef] +min = 0.00001 +max = 0.01 + +[sweep.train.gamma] +min = 0.88 +max = 0.99 + +[sweep.train.gae_lambda] +min = 0.8 +max = 0.95 + +[sweep.train.clip_coef] +min = 0.1 +max = 0.35 + +[sweep.train.vf_coef] +min = 0.5 +max = 2.0 + +[sweep.train.vf_clip_coef] +min = 1.0 +max = 5.0 diff --git a/ocean/mujoco/backend.h b/ocean/mujoco/backend.h index 2d3aab72ae..e7c28663bf 100644 --- a/ocean/mujoco/backend.h +++ b/ocean/mujoco/backend.h @@ -1,13 +1,4 @@ -// Vector backends for the ocean/mujoco envs, included at the end of each -// mjc_.h after Env, mjc_reset, mjc_step, mjc_render and mjc_init. All -// memory is bound at create time: the solver scratch (MJ_SCRATCH floats per -// env, see mj_makeData) and, on the GPU, the whole batch. CPU: the puf_* per-env -// API calls straight through. GPU (mjc_.cu defines PUF_BACKEND PUF_GPU): -// one thread per env steps a thread-local Env (local memory is lane interleaved, -// so every access coalesces) and copies only the persistent prefix, everything -// before MjData.xquat, in and out of the device batch; the scratch is laid out -// lane interleaved too (element stride 32). The trainer's device obs/action/ -// reward/terminal buffers are bound into Env.agents at create time. +// Vector backend, included at the end of each mjc_.h. #if PUF_BACKEND == PUF_GPU #define MJC_BLOCK 128 @@ -19,24 +10,25 @@ struct { cudaStream_t stream; } mjc_gpu; -__global__ void mjc_reset_kernel(Env* envs, int n) { +__global__ void mjc_kernel(Env* envs, int n, int reset) { int i = blockIdx.x*blockDim.x + threadIdx.x; if (i < n) { Env env; memcpy(&env, &envs[i], MJC_STATE); - mjc_reset(&env); + if (reset) { + mjc_reset(&env); + } else { + mjc_step(&env); + } memcpy(&envs[i], &env, MJC_STATE); } } -__global__ void mjc_step_kernel(Env* envs, int n) { - int i = blockIdx.x*blockDim.x + threadIdx.x; - if (i < n) { - Env env; - memcpy(&env, &envs[i], MJC_STATE); - mjc_step(&env); - memcpy(&envs[i], &env, MJC_STATE); - } +// The trainer resets before binding its stream, so resets go to the null stream +void mjc_launch(int reset) { + mjc_kernel<<<(mjc_gpu.n + MJC_BLOCK - 1) / MJC_BLOCK, MJC_BLOCK, 0, mjc_gpu.stream>>>( + mjc_gpu.envs, mjc_gpu.n, reset); + assert(cudaGetLastError() == cudaSuccess); } Env* puf_vec_create(int n, Dict* kwargs, obs_t* observations, float* actions, float* rewards, @@ -45,8 +37,7 @@ Env* puf_vec_create(int n, Dict* kwargs, obs_t* observations, float* actions, fl MjModel* m; assert(cudaMalloc((void**)&m, sizeof(MjModel)) == cudaSuccess); float* scratch; - size_t groups = (n + 31) / 32; - assert(cudaMalloc((void**)&scratch, groups*32*MJ_SCRATCH*sizeof(float)) == cudaSuccess + assert(cudaMalloc((void**)&scratch, sizeof(float)*((n + 31) / 32*32)*MJ_SCRATCH) == cudaSuccess && "GPU env solver scratch does not fit in device memory"); for (int i = 0; i < n; i++) { Env* env = &host[i]; @@ -76,19 +67,14 @@ void puf_init(Env* env, Dict* kwargs) { } void puf_reset(Env* envs) { - mjc_reset_kernel<<<(mjc_gpu.n + MJC_BLOCK - 1) / MJC_BLOCK, MJC_BLOCK>>>(mjc_gpu.envs, - mjc_gpu.n); - assert(cudaGetLastError() == cudaSuccess); + mjc_launch(1); } void puf_step(Env* envs) { - mjc_step_kernel<<<(mjc_gpu.n + MJC_BLOCK - 1) / MJC_BLOCK, MJC_BLOCK, 0, mjc_gpu.stream>>>( - mjc_gpu.envs, mjc_gpu.n); - assert(cudaGetLastError() == cudaSuccess); + mjc_launch(0); } -// Copy env 0's state back and draw it with the host model (mj_render recomputes -// kinematics and contacts from qpos) +// Draw env 0 with the host model (mj_render recomputes kinematics from qpos) void puf_render(Env* envs) { Env env; cudaStreamSynchronize(mjc_gpu.stream); @@ -109,6 +95,40 @@ void puf_init(Env* env, Dict* kwargs) { mj_makeData(&env->d, (float*)calloc(MJ_SCRATCH, sizeof(float)), 1); } +#ifdef PUFFERCPU_EVAL_MAIN +// Whole control steps (MJC_FRAME_SKIP substeps at once) strobe on screen at +// 20 Hz. Run the real mjc_step for exact rewards, logs and resets, then +// rewind and re-integrate the same deterministic substeps one per rendered +// frame; PUF_STEPS_PER_SEC is therefore the substep rate 1/opt_timestep. +#define PUF_EVAL_SHOULD_FORWARD +#define MJC_STATE (offsetof(MjData, xquat)) +char mjc_true[MJC_STATE]; + +void puf_reset(Env* env) { + mjc_reset(env); + memcpy(mjc_true, &env->d, MJC_STATE); +} + +void puf_step(Env* env) { + MjData* d = &env->d; + if (env->tick_frames_left > 0) { + env->tick_frames_left--; + mj_step(env->m, d); + return; + } + memcpy(d, mjc_true, MJC_STATE); + mjc_step(env); + float ctrl[MJ_MAX_NU]; + memcpy(ctrl, d->ctrl, sizeof(ctrl)); + char post[MJC_STATE]; + memcpy(post, d, MJC_STATE); + memcpy(d, mjc_true, MJC_STATE); + memcpy(mjc_true, post, MJC_STATE); + memcpy(d->ctrl, ctrl, sizeof(ctrl)); + mj_step(env->m, d); + env->tick_frames_left = MJC_FRAME_SKIP - 1; +} +#else void puf_reset(Env* env) { mjc_reset(env); } @@ -116,6 +136,7 @@ void puf_reset(Env* env) { void puf_step(Env* env) { mjc_step(env); } +#endif void puf_render(Env* env) { mjc_render(env); diff --git a/ocean/mujoco/mjc_ant.h b/ocean/mujoco/mjc_ant.h index a0bc7e5261..fd9380bd0d 100644 --- a/ocean/mujoco/mjc_ant.h +++ b/ocean/mujoco/mjc_ant.h @@ -1,8 +1,6 @@ -// Ant (gymnasium Ant-v5) on the MuJoCo-style physics core: obs = qpos[2:] + -// qvel + clip(cfrc_ext[1:], -1, 1), reward = healthy + forward velocity - ctrl -// cost - contact cost, terminates when the torso leaves [0.2, 1.0] m. Model: -// resources/mujoco/ant.xml compiled by mjcf2bin.py. mjc_reset/mjc_step are -// host+device (see backend.h). +// Ant (gymnasium Ant-v5): obs = qpos[2:] + qvel scaled by mj_obsJoints + +// clip(cfrc_ext[1:], -1, 1), reward = healthy + forward velocity - ctrl cost +// - contact cost, terminates when the torso leaves z in [0.2, 1.0]. #include #include @@ -11,7 +9,7 @@ #include "raylib.h" typedef float obs_t; #include "pufferenv.h" -// physics.h capacities sized to this model (checked by mj_loadModel) +// model capacities, checked by mj_loadModel #define MJ_MAX_NQ 15 #define MJ_MAX_NV 14 #define MJ_MAX_NBODY 14 @@ -23,11 +21,13 @@ typedef float obs_t; #include "physics.h" #include "render.h" -#define ANT_FRAME_SKIP 5 +// forward velocity worth perf 1 and a trainer reward of 1 per step +#define ANT_TARGET_VEL 5.0f +#define MJC_FRAME_SKIP 5 #define OBS_SIZE 105 #define NUM_ATNS 8 #define ACT_SIZES {1, 1, 1, 1, 1, 1, 1, 1} -#define PUF_STEPS_PER_SEC 20 +#define PUF_STEPS_PER_SEC 100 MjModel mj_model; @@ -50,6 +50,7 @@ struct Env { unsigned int rng; const MjModel* m; int tick; + int tick_frames_left; float x_start; float episode_return; int max_steps; @@ -65,10 +66,7 @@ typedef Env Ant; // Returns the contact cost (sum of squared clipped external forces) MJ_HD float compute_observations(Ant* env) { const MjModel* m = env->m; - float* obs = env->agents[0].observations; - memcpy(obs, env->d.qpos + 2, (m->nq - 2)*sizeof(float)); - memcpy(obs + m->nq - 2, env->d.qvel, m->nv*sizeof(float)); - float* cfrc = obs + m->nq - 2 + m->nv; + float* cfrc = mj_obsJoints(m, &env->d, env->agents[0].observations, 2, 20.0f); float cost = 0.0f; for (int b = 1; b < m->nbody; b++) { for (int k = 0; k < 6; k++) { @@ -105,15 +103,14 @@ MJ_HD void mjc_step(Ant* env) { env->d.ctrl[i] = fminf(fmaxf(actions[i], -1.0f), 1.0f); cost += env->d.ctrl[i]*env->d.ctrl[i]; } - // Gym measures the torso displacement with body xpos, which lags qpos by - // one substep (kinematics of the last forward pass) + // Gym measures displacement with body xpos, which lags qpos by one substep float x0 = env->d.xpos[1][0]; - for (int k = 0; k < ANT_FRAME_SKIP; k++) { + for (int k = 0; k < MJC_FRAME_SKIP; k++) { mj_step(m, &env->d); } mj_rnePostConstraint(m, &env->d); float* q = env->d.qpos; - float dt = ANT_FRAME_SKIP*m->opt_timestep; + float dt = MJC_FRAME_SKIP*m->opt_timestep; float x_velocity = (env->d.xpos[1][0] - x0) / dt; int healthy = isfinite(q[2]) && q[2] >= 0.2f && q[2] <= 1.0f; float contact_cost = compute_observations(env); @@ -121,7 +118,7 @@ MJ_HD void mjc_step(Ant* env) { - contact_cost + (healthy ? env->healthy_reward : 0.0f); env->tick++; env->episode_return += reward; - env->agents[0].rewards[0] = reward; + env->agents[0].rewards[0] = reward / ANT_TARGET_VEL; env->agents[0].terminals[0] = 0.0f; if (healthy && env->tick < env->max_steps) { return; @@ -129,7 +126,7 @@ MJ_HD void mjc_step(Ant* env) { float distance = q[0] - env->x_start; float xvel = distance / (env->tick*dt); env->agents[0].terminals[0] = 1.0f; - env->log.perf += fminf(fmaxf(xvel / 5.0f, 0.0f), 1.0f); + env->log.perf += fminf(fmaxf(xvel / ANT_TARGET_VEL, 0.0f), 1.0f); env->log.score += env->episode_return; env->log.episode_return += env->episode_return; env->log.episode_length += env->tick; @@ -155,7 +152,6 @@ void mjc_init(Ant* env, Dict* kwargs) { env->m = &mj_model; env->num_agents = 1; env->agents[0].policy = 0; - env->agents[0].action_mask = NULL; env->max_steps = dict_get(kwargs, "max_steps"); env->reset_noise_scale = dict_get(kwargs, "reset_noise_scale"); env->forward_reward_weight = dict_get(kwargs, "forward_reward_weight"); diff --git a/ocean/mujoco/mjc_half_cheetah.h b/ocean/mujoco/mjc_half_cheetah.h index 158f912d74..95c9973b1b 100644 --- a/ocean/mujoco/mjc_half_cheetah.h +++ b/ocean/mujoco/mjc_half_cheetah.h @@ -1,8 +1,5 @@ -// HalfCheetah (gymnasium HalfCheetah-v5) on the MuJoCo-style physics core: -// obs = qpos[1:] + qvel, reward = forward velocity - ctrl cost, 1000 step -// episodes. Model: resources/mujoco/half_cheetah.xml compiled by mjcf2bin.py. -// mjc_reset/mjc_step are host+device; backend.h wraps them for the CPU vec -// or launches them one thread per env on the GPU (--cu). +// HalfCheetah (gymnasium HalfCheetah-v5): obs = qpos[1:] + qvel scaled by +// mj_obsJoints, reward = forward velocity - ctrl cost, 1000 step episodes. #include #include @@ -11,7 +8,7 @@ #include "raylib.h" typedef float obs_t; #include "pufferenv.h" -// physics.h capacities sized to this model (checked by mj_loadModel) +// model capacities, checked by mj_loadModel #define MJ_MAX_NQ 9 #define MJ_MAX_NV 9 #define MJ_MAX_NBODY 8 @@ -23,11 +20,13 @@ typedef float obs_t; #include "physics.h" #include "render.h" -#define HC_FRAME_SKIP 5 +// forward velocity worth perf 1 and a trainer reward of 1 per step +#define HC_TARGET_VEL 10.0f +#define MJC_FRAME_SKIP 5 #define OBS_SIZE 17 #define NUM_ATNS 6 #define ACT_SIZES {1, 1, 1, 1, 1, 1} -#define PUF_STEPS_PER_SEC 20 +#define PUF_STEPS_PER_SEC 100 MjModel mj_model; @@ -51,6 +50,7 @@ struct Env { unsigned int rng; const MjModel* m; int tick; + int tick_frames_left; float x_start; float episode_return; float episode_ctrl_cost; @@ -62,12 +62,6 @@ struct Env { }; typedef Env HalfCheetah; -MJ_HD void compute_observations(HalfCheetah* env) { - float* obs = env->agents[0].observations; - memcpy(obs, env->d.qpos + 1, (env->m->nq - 1)*sizeof(float)); - memcpy(obs + env->m->nq - 1, env->d.qvel, env->m->nv*sizeof(float)); -} - MJ_HD void mjc_reset(HalfCheetah* env) { const MjModel* m = env->m; mj_resetData(m, &env->d); @@ -82,7 +76,7 @@ MJ_HD void mjc_reset(HalfCheetah* env) { env->x_start = env->d.qpos[0]; env->episode_return = 0.0f; env->episode_ctrl_cost = 0.0f; - compute_observations(env); + mj_obsJoints(m, &env->d, env->agents[0].observations, 1, 20.0f); } MJ_HD void mjc_step(HalfCheetah* env) { @@ -94,25 +88,25 @@ MJ_HD void mjc_step(HalfCheetah* env) { cost += env->d.ctrl[i]*env->d.ctrl[i]; } float x0 = env->d.qpos[0]; - for (int k = 0; k < HC_FRAME_SKIP; k++) { + for (int k = 0; k < MJC_FRAME_SKIP; k++) { mj_step(m, &env->d); } - float dt = HC_FRAME_SKIP*m->opt_timestep; + float dt = MJC_FRAME_SKIP*m->opt_timestep; float x_velocity = (env->d.qpos[0] - x0) / dt; float reward = env->forward_reward_weight*x_velocity - env->ctrl_cost_weight*cost; env->tick++; env->episode_return += reward; env->episode_ctrl_cost += env->ctrl_cost_weight*cost; - env->agents[0].rewards[0] = reward; + env->agents[0].rewards[0] = reward / HC_TARGET_VEL; env->agents[0].terminals[0] = 0.0f; if (env->tick < env->max_steps) { - compute_observations(env); + mj_obsJoints(m, &env->d, env->agents[0].observations, 1, 20.0f); return; } float distance = env->d.qpos[0] - env->x_start; float xvel = distance / (env->tick*dt); env->agents[0].terminals[0] = 1.0f; - env->log.perf += fminf(fmaxf(xvel / 10.0f, 0.0f), 1.0f); + env->log.perf += fminf(fmaxf(xvel / HC_TARGET_VEL, 0.0f), 1.0f); env->log.score += env->episode_return; env->log.episode_return += env->episode_return; env->log.episode_length += env->tick; @@ -138,7 +132,6 @@ void mjc_init(HalfCheetah* env, Dict* kwargs) { env->m = &mj_model; env->num_agents = 1; env->agents[0].policy = 0; - env->agents[0].action_mask = NULL; env->max_steps = dict_get(kwargs, "max_steps"); env->reset_noise_scale = dict_get(kwargs, "reset_noise_scale"); env->forward_reward_weight = dict_get(kwargs, "forward_reward_weight"); diff --git a/ocean/mujoco/mjc_hopper.h b/ocean/mujoco/mjc_hopper.h index 3eff17b746..594333610c 100644 --- a/ocean/mujoco/mjc_hopper.h +++ b/ocean/mujoco/mjc_hopper.h @@ -1,7 +1,5 @@ -// Hopper (gymnasium Hopper-v5) on the MuJoCo-style physics core: obs = -// qpos[1:] + clip(qvel, -10, 10), reward = healthy + forward velocity - ctrl -// cost, terminates when unhealthy. Model: resources/mujoco/hopper.xml compiled -// by mjcf2bin.py. mjc_reset/mjc_step are host+device (see backend.h). +// Hopper (gymnasium Hopper-v5): obs = qpos[1:] + qvel scaled by mj_obsJoints, +// reward = healthy + forward velocity - ctrl cost, terminates when unhealthy. #include #include @@ -10,7 +8,7 @@ #include "raylib.h" typedef float obs_t; #include "pufferenv.h" -// physics.h capacities sized to this model +// model capacities, checked by mj_loadModel #define MJ_MAX_NQ 6 #define MJ_MAX_NV 6 #define MJ_MAX_NBODY 5 @@ -22,11 +20,13 @@ typedef float obs_t; #include "physics.h" #include "render.h" -#define HP_FRAME_SKIP 4 +// forward velocity worth perf 1 and a trainer reward of 1 per step +#define HP_TARGET_VEL 3.0f +#define MJC_FRAME_SKIP 4 #define OBS_SIZE 11 #define NUM_ATNS 3 #define ACT_SIZES {1, 1, 1} -#define PUF_STEPS_PER_SEC 125 +#define PUF_STEPS_PER_SEC 500 MjModel mj_model; @@ -49,6 +49,7 @@ struct Env { unsigned int rng; const MjModel* m; int tick; + int tick_frames_left; float x_start; float episode_return; int max_steps; @@ -60,15 +61,6 @@ struct Env { }; typedef Env Hopper; -MJ_HD void compute_observations(Hopper* env) { - const MjModel* m = env->m; - float* obs = env->agents[0].observations; - memcpy(obs, env->d.qpos + 1, (m->nq - 1)*sizeof(float)); - for (int i = 0; i < m->nv; i++) { - obs[m->nq - 1 + i] = fminf(fmaxf(env->d.qvel[i], -10.0f), 10.0f); - } -} - MJ_HD void mjc_reset(Hopper* env) { const MjModel* m = env->m; mj_resetData(m, &env->d); @@ -82,7 +74,7 @@ MJ_HD void mjc_reset(Hopper* env) { env->tick = 0; env->x_start = env->d.qpos[0]; env->episode_return = 0.0f; - compute_observations(env); + mj_obsJoints(m, &env->d, env->agents[0].observations, 1, 10.0f); } MJ_HD void mjc_step(Hopper* env) { @@ -94,11 +86,11 @@ MJ_HD void mjc_step(Hopper* env) { cost += env->d.ctrl[i]*env->d.ctrl[i]; } float x0 = env->d.qpos[0]; - for (int k = 0; k < HP_FRAME_SKIP; k++) { + for (int k = 0; k < MJC_FRAME_SKIP; k++) { mj_step(m, &env->d); } float* q = env->d.qpos; - float dt = HP_FRAME_SKIP*m->opt_timestep; + float dt = MJC_FRAME_SKIP*m->opt_timestep; float x_velocity = (q[0] - x0) / dt; int healthy = q[1] > 0.7f && q[2] > -0.2f && q[2] < 0.2f; for (int i = 2; i < m->nq; i++) { @@ -111,16 +103,16 @@ MJ_HD void mjc_step(Hopper* env) { + (healthy ? env->healthy_reward : 0.0f); env->tick++; env->episode_return += reward; - env->agents[0].rewards[0] = reward; + env->agents[0].rewards[0] = reward / HP_TARGET_VEL; env->agents[0].terminals[0] = 0.0f; if (healthy && env->tick < env->max_steps) { - compute_observations(env); + mj_obsJoints(m, &env->d, env->agents[0].observations, 1, 10.0f); return; } float distance = q[0] - env->x_start; float xvel = distance / (env->tick*dt); env->agents[0].terminals[0] = 1.0f; - env->log.perf += fminf(fmaxf(xvel / 3.0f, 0.0f), 1.0f); + env->log.perf += fminf(fmaxf(xvel / HP_TARGET_VEL, 0.0f), 1.0f); env->log.score += env->episode_return; env->log.episode_return += env->episode_return; env->log.episode_length += env->tick; @@ -145,7 +137,6 @@ void mjc_init(Hopper* env, Dict* kwargs) { env->m = &mj_model; env->num_agents = 1; env->agents[0].policy = 0; - env->agents[0].action_mask = NULL; env->max_steps = dict_get(kwargs, "max_steps"); env->reset_noise_scale = dict_get(kwargs, "reset_noise_scale"); env->forward_reward_weight = dict_get(kwargs, "forward_reward_weight"); diff --git a/ocean/mujoco/mjc_humanoid.h b/ocean/mujoco/mjc_humanoid.h index 0d3e9eaff1..6fa6facaf7 100644 --- a/ocean/mujoco/mjc_humanoid.h +++ b/ocean/mujoco/mjc_humanoid.h @@ -1,9 +1,7 @@ -// Humanoid (gymnasium Humanoid-v5) on the MuJoCo-style physics core: obs = -// qpos[2:] + qvel + cinert[1:] + cvel[1:] + qfrc_actuator[6:] + cfrc_ext[1:], -// reward = healthy + forward COM velocity - ctrl cost - contact cost (clamped -// at 10), terminates when the torso leaves z in (1, 2). Model: -// resources/mujoco/humanoid.xml compiled by mjcf2bin.py. mjc_reset/mjc_step -// are host+device (see backend.h). +// Humanoid (gymnasium Humanoid-v5): obs = qpos[2:] + qvel + cinert[1:] + +// cvel[1:] + qfrc_actuator[6:] + cfrc_ext[1:], each block scaled to O(1), +// reward = healthy + forward COM velocity - ctrl cost - contact cost, +// terminates when the torso leaves z in (1, 2). #include #include @@ -12,7 +10,7 @@ #include "raylib.h" typedef float obs_t; #include "pufferenv.h" -// physics.h capacities sized to this model (checked by mj_loadModel) +// model capacities, checked by mj_loadModel #define MJ_MAX_NQ 24 #define MJ_MAX_NV 23 #define MJ_MAX_NBODY 14 @@ -24,12 +22,17 @@ typedef float obs_t; #include "physics.h" #include "render.h" -#define HM_FRAME_SKIP 5 +// forward velocity worth perf 1 and a trainer reward of 1 per step +#define HM_TARGET_VEL 3.0f +#define MJC_FRAME_SKIP 5 #define HM_CONTACT_COST_MAX 10.0f +#define HM_VEL_SCALE 20.0f +#define HM_INERTIA_SCALE 10.0f +#define HM_FORCE_SCALE 100.0f #define OBS_SIZE 348 #define NUM_ATNS 17 #define ACT_SIZES {1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1} -#define PUF_STEPS_PER_SEC 67 +#define PUF_STEPS_PER_SEC 333 MjModel mj_model; @@ -52,6 +55,7 @@ struct Env { unsigned int rng; const MjModel* m; int tick; + int tick_frames_left; float x_start; float episode_return; int max_steps; @@ -67,18 +71,11 @@ typedef Env Humanoid; // Returns the contact cost (sum of squared external forces, clamped) MJ_HD float compute_observations(Humanoid* env) { const MjModel* m = env->m; - float* obs = env->agents[0].observations; - memcpy(obs, env->d.qpos + 2, (m->nq - 2)*sizeof(float)); - obs += m->nq - 2; - memcpy(obs, env->d.qvel, m->nv*sizeof(float)); - obs += m->nv; - memcpy(obs, env->d.cinert[1], 10*(m->nbody - 1)*sizeof(float)); - obs += 10*(m->nbody - 1); - memcpy(obs, env->d.cvel[1], 6*(m->nbody - 1)*sizeof(float)); - obs += 6*(m->nbody - 1); - memcpy(obs, env->d.qfrc_actuator + 6, (m->nv - 6)*sizeof(float)); - obs += m->nv - 6; - memcpy(obs, env->d.cfrc_ext[1], 6*(m->nbody - 1)*sizeof(float)); + float* obs = mj_obsJoints(m, &env->d, env->agents[0].observations, 2, HM_VEL_SCALE); + obs = mj_obsScaled(obs, env->d.cinert[1], 10*(m->nbody - 1), HM_INERTIA_SCALE); + obs = mj_obsScaled(obs, env->d.cvel[1], 6*(m->nbody - 1), HM_VEL_SCALE); + obs = mj_obsScaled(obs, env->d.qfrc_actuator + 6, m->nv - 6, HM_FORCE_SCALE); + mj_obsScaled(obs, env->d.cfrc_ext[1], 6*(m->nbody - 1), HM_FORCE_SCALE); float cost = 0.0f; for (int b = 1; b < m->nbody; b++) { cost += mju_dot6(env->d.cfrc_ext[b], env->d.cfrc_ext[b]); @@ -86,8 +83,7 @@ MJ_HD float compute_observations(Humanoid* env) { return fminf(env->contact_cost_weight*cost, HM_CONTACT_COST_MAX); } -// x of the whole-body center of mass from the inertial frames of the last -// kinematics pass (Gym's mass_center) +// Gym's mass_center: whole-body COM x from the last kinematics pass MJ_HD float mjc_com_x(Humanoid* env) { const MjModel* m = env->m; float num = 0.0f; @@ -126,12 +122,12 @@ MJ_HD void mjc_step(Humanoid* env) { cost += env->d.ctrl[i]*env->d.ctrl[i]; } float x0 = mjc_com_x(env); - for (int k = 0; k < HM_FRAME_SKIP; k++) { + for (int k = 0; k < MJC_FRAME_SKIP; k++) { mj_step(m, &env->d); } mj_rnePostConstraint(m, &env->d); float* q = env->d.qpos; - float dt = HM_FRAME_SKIP*m->opt_timestep; + float dt = MJC_FRAME_SKIP*m->opt_timestep; float x_velocity = (mjc_com_x(env) - x0) / dt; int healthy = q[2] > 1.0f && q[2] < 2.0f; float contact_cost = compute_observations(env); @@ -139,7 +135,7 @@ MJ_HD void mjc_step(Humanoid* env) { - contact_cost + (healthy ? env->healthy_reward : 0.0f); env->tick++; env->episode_return += reward; - env->agents[0].rewards[0] = reward; + env->agents[0].rewards[0] = reward / HM_TARGET_VEL; env->agents[0].terminals[0] = 0.0f; if (healthy && env->tick < env->max_steps) { return; @@ -147,7 +143,7 @@ MJ_HD void mjc_step(Humanoid* env) { float distance = q[0] - env->x_start; float xvel = distance / (env->tick*dt); env->agents[0].terminals[0] = 1.0f; - env->log.perf += fminf(fmaxf(xvel / 3.0f, 0.0f), 1.0f); + env->log.perf += fminf(fmaxf(xvel / HM_TARGET_VEL, 0.0f), 1.0f); env->log.score += env->episode_return; env->log.episode_return += env->episode_return; env->log.episode_length += env->tick; @@ -173,7 +169,6 @@ void mjc_init(Humanoid* env, Dict* kwargs) { env->m = &mj_model; env->num_agents = 1; env->agents[0].policy = 0; - env->agents[0].action_mask = NULL; env->max_steps = dict_get(kwargs, "max_steps"); env->reset_noise_scale = dict_get(kwargs, "reset_noise_scale"); env->forward_reward_weight = dict_get(kwargs, "forward_reward_weight"); diff --git a/ocean/mujoco/mjc_walker2d.h b/ocean/mujoco/mjc_walker2d.h index 24386cc221..57f98426bc 100644 --- a/ocean/mujoco/mjc_walker2d.h +++ b/ocean/mujoco/mjc_walker2d.h @@ -1,8 +1,6 @@ -// Walker2d (gymnasium Walker2d-v5) on the MuJoCo-style physics core: obs = -// qpos[1:] + clip(qvel, -10, 10), reward = healthy + forward velocity - ctrl -// cost, terminates when the torso leaves z in (0.8, 2) or |angle| >= 1. -// Model: resources/mujoco/walker2d.xml (gymnasium walker2d_v5.xml) compiled by -// mjcf2bin.py. mjc_reset/mjc_step are host+device (see backend.h). +// Walker2d (gymnasium Walker2d-v5): obs = qpos[1:] + qvel scaled by +// mj_obsJoints, reward = healthy + forward velocity - ctrl cost, terminates +// when the torso leaves z in (0.8, 2) or |angle| >= 1. #include #include @@ -11,7 +9,7 @@ #include "raylib.h" typedef float obs_t; #include "pufferenv.h" -// physics.h capacities sized to this model (checked by mj_loadModel) +// model capacities, checked by mj_loadModel #define MJ_MAX_NQ 9 #define MJ_MAX_NV 9 #define MJ_MAX_NBODY 8 @@ -23,11 +21,13 @@ typedef float obs_t; #include "physics.h" #include "render.h" -#define WK_FRAME_SKIP 4 +// forward velocity worth perf 1 and a trainer reward of 1 per step +#define WK_TARGET_VEL 5.0f +#define MJC_FRAME_SKIP 4 #define OBS_SIZE 17 #define NUM_ATNS 6 #define ACT_SIZES {1, 1, 1, 1, 1, 1} -#define PUF_STEPS_PER_SEC 125 +#define PUF_STEPS_PER_SEC 500 MjModel mj_model; @@ -50,6 +50,7 @@ struct Env { unsigned int rng; const MjModel* m; int tick; + int tick_frames_left; float x_start; float episode_return; int max_steps; @@ -61,15 +62,6 @@ struct Env { }; typedef Env Walker2d; -MJ_HD void compute_observations(Walker2d* env) { - const MjModel* m = env->m; - float* obs = env->agents[0].observations; - memcpy(obs, env->d.qpos + 1, (m->nq - 1)*sizeof(float)); - for (int i = 0; i < m->nv; i++) { - obs[m->nq - 1 + i] = fminf(fmaxf(env->d.qvel[i], -10.0f), 10.0f); - } -} - MJ_HD void mjc_reset(Walker2d* env) { const MjModel* m = env->m; mj_resetData(m, &env->d); @@ -83,7 +75,7 @@ MJ_HD void mjc_reset(Walker2d* env) { env->tick = 0; env->x_start = env->d.qpos[0]; env->episode_return = 0.0f; - compute_observations(env); + mj_obsJoints(m, &env->d, env->agents[0].observations, 1, 10.0f); } MJ_HD void mjc_step(Walker2d* env) { @@ -95,27 +87,27 @@ MJ_HD void mjc_step(Walker2d* env) { cost += env->d.ctrl[i]*env->d.ctrl[i]; } float x0 = env->d.qpos[0]; - for (int k = 0; k < WK_FRAME_SKIP; k++) { + for (int k = 0; k < MJC_FRAME_SKIP; k++) { mj_step(m, &env->d); } float* q = env->d.qpos; - float dt = WK_FRAME_SKIP*m->opt_timestep; + float dt = MJC_FRAME_SKIP*m->opt_timestep; float x_velocity = (q[0] - x0) / dt; int healthy = q[1] > 0.8f && q[1] < 2.0f && q[2] > -1.0f && q[2] < 1.0f; float reward = env->forward_reward_weight*x_velocity - env->ctrl_cost_weight*cost + (healthy ? env->healthy_reward : 0.0f); env->tick++; env->episode_return += reward; - env->agents[0].rewards[0] = reward; + env->agents[0].rewards[0] = reward / WK_TARGET_VEL; env->agents[0].terminals[0] = 0.0f; if (healthy && env->tick < env->max_steps) { - compute_observations(env); + mj_obsJoints(m, &env->d, env->agents[0].observations, 1, 10.0f); return; } float distance = q[0] - env->x_start; float xvel = distance / (env->tick*dt); env->agents[0].terminals[0] = 1.0f; - env->log.perf += fminf(fmaxf(xvel / 5.0f, 0.0f), 1.0f); + env->log.perf += fminf(fmaxf(xvel / WK_TARGET_VEL, 0.0f), 1.0f); env->log.score += env->episode_return; env->log.episode_return += env->episode_return; env->log.episode_length += env->tick; @@ -140,7 +132,6 @@ void mjc_init(Walker2d* env, Dict* kwargs) { env->m = &mj_model; env->num_agents = 1; env->agents[0].policy = 0; - env->agents[0].action_mask = NULL; env->max_steps = dict_get(kwargs, "max_steps"); env->reset_noise_scale = dict_get(kwargs, "reset_noise_scale"); env->forward_reward_weight = dict_get(kwargs, "forward_reward_weight"); diff --git a/ocean/mujoco/mjcf2bin.py b/ocean/mujoco/mjcf2bin.py index 84ef892cf3..d659538112 100644 --- a/ocean/mujoco/mjcf2bin.py +++ b/ocean/mujoco/mjcf2bin.py @@ -23,6 +23,10 @@ def compile_model(xml, out): assert all(m.actuator_trntype == 0) and all(m.actuator_dyntype == 0), "motors only" assert all(m.jnt_type[m.actuator_trnid[:, 0]] >= 2), "actuators on hinge/slide joints only" assert m.opt.cone == 0, "pyramidal cones only" + assert m.opt.integrator <= 1, "Euler/RK4 integrators only" + assert all(np.isin(m.geom_type, [0, 2, 3])), "plane/sphere/capsule geoms only" + assert all(m.jnt_type[m.jnt_limited.astype(bool) | (m.jnt_stiffness != 0)] >= 2), \ + "limits and springs on hinge/slide joints only" # fixed tendons without limits, springs or dampers (humanoid) exert no force assert not m.tendon_limited.any() and not m.tendon_stiffness.any() \ and not m.tendon_damping.any(), "tendon limits/springs/dampers unsupported" diff --git a/ocean/mujoco/physics.h b/ocean/mujoco/physics.h index addb9c13e9..e41ae9ded1 100644 --- a/ocean/mujoco/physics.h +++ b/ocean/mujoco/physics.h @@ -1,12 +1,4 @@ -// MuJoCo-style rigid body physics for PufferLib envs. CUDA C99, fixed capacity, -// float. Ports mj_step's pipeline (same algorithms, names and layouts as the -// MuJoCo engine) for the features the classic Gym models use: free/ball/hinge/ -// slide joints, plane/sphere/capsule geoms, motor actuators, joint springs and -// dampers, joint limits and pyramidal friction contacts with MuJoCo's soft -// constraint model, semi-implicit Euler with implicit damping, and RK4. -// Models are compiled from MJCF by ocean/mujoco/mjcf2bin.py and loaded with -// mj_loadModel; arrays have fixed MJ_MAX_* capacity (defaults below, envs -// define tighter ones for their model before including) and runtime counts. +// MuJoCo-style rigid body physics for PufferLib envs. #include #include @@ -51,7 +43,7 @@ #ifndef MJ_MAXEFC #define MJ_MAXEFC 128 #endif -#define MJ_MAGIC 0x4e424a4d +#define MJ_MAGIC 0x4e424a4d // Basically a checksum to make sure the loaded model is valid enum {MJ_JNT_FREE, MJ_JNT_BALL, MJ_JNT_SLIDE, MJ_JNT_HINGE}; enum {MJ_GEOM_PLANE, MJ_GEOM_HFIELD, MJ_GEOM_SPHERE, MJ_GEOM_CAPSULE, MJ_GEOM_ELLIPSOID, MJ_GEOM_CYLINDER, MJ_GEOM_BOX}; @@ -134,7 +126,6 @@ typedef struct { int efc_address; } MjContact; -// Constraint solver scratch: MJ_SCRATCH floats per env #define MJ_SCRATCH (2*MJ_MAXEFC*MJ_MAXEFC + MJ_MAXEFC*MJ_MAX_NV) typedef struct { @@ -200,16 +191,10 @@ void mj_read(FILE* fp, void* dst, int count, int size) { void mj_loadModel(MjModel* m, const char* path) { FILE* fp = fopen(path, "rb"); assert(fp && "cannot open model file"); - int head[9]; - mj_read(fp, head, 9, sizeof(int)); + int head[2]; + mj_read(fp, head, 2, sizeof(int)); assert(head[0] == MJ_MAGIC && head[1] == 1 && "bad model file"); - m->nq = head[2]; - m->nv = head[3]; - m->nbody = head[4]; - m->njnt = head[5]; - m->ngeom = head[6]; - m->nsite = head[7]; - m->nu = head[8]; + mj_read(fp, &m->nq, 7, sizeof(int)); assert(m->nq <= MJ_MAX_NQ && m->nv <= MJ_MAX_NV && m->nbody <= MJ_MAX_NBODY && m->njnt <= MJ_MAX_NJNT && m->ngeom <= MJ_MAX_NGEOM && m->nsite <= MJ_MAX_NSITE && m->nu <= MJ_MAX_NU && "model exceeds MJ_MAX_* capacity"); @@ -321,11 +306,16 @@ MJ_HD void mju_normalize4(float* q) { } } +// res may alias qa (kinematics accumulates joint rotations in place) MJ_HD void mju_mulQuat(float* res, const float* qa, const float* qb) { - res[0] = qa[0]*qb[0] - qa[1]*qb[1] - qa[2]*qb[2] - qa[3]*qb[3]; - res[1] = qa[0]*qb[1] + qa[1]*qb[0] + qa[2]*qb[3] - qa[3]*qb[2]; - res[2] = qa[0]*qb[2] - qa[1]*qb[3] + qa[2]*qb[0] + qa[3]*qb[1]; - res[3] = qa[0]*qb[3] + qa[1]*qb[2] - qa[2]*qb[1] + qa[3]*qb[0]; + float t0 = qa[0]*qb[0] - qa[1]*qb[1] - qa[2]*qb[2] - qa[3]*qb[3]; + float t1 = qa[0]*qb[1] + qa[1]*qb[0] + qa[2]*qb[3] - qa[3]*qb[2]; + float t2 = qa[0]*qb[2] - qa[1]*qb[3] + qa[2]*qb[0] + qa[3]*qb[1]; + float t3 = qa[0]*qb[3] + qa[1]*qb[2] - qa[2]*qb[1] + qa[3]*qb[0]; + res[0] = t0; + res[1] = t1; + res[2] = t2; + res[3] = t3; } MJ_HD void mju_rotVecQuat(float* res, const float* vec, const float* quat) { @@ -449,6 +439,42 @@ MJ_HD float mju_randn(unsigned int* rng) { return sqrtf(-2.0f*logf(u1))*cosf(2.0f*(float)M_PI*u2); } +// Gym joint obs (qpos[skip:] + qvel) scaled to O(1): limited joints by range, +// unlimited hinges by pi, qvel by vel_scale, all clamped to [-1, 1]; quats and +// unlimited slides as is. Returns the advanced obs pointer. +MJ_HD float* mj_obsJoints(const MjModel* m, const MjData* d, float* obs, int skip, + float vel_scale) { + float q[MJ_MAX_NQ]; + for (int j = 0; j < m->njnt; j++) { + int adr = m->jnt_qposadr[j]; + int type = m->jnt_type[j]; + if (type == MJ_JNT_FREE || type == MJ_JNT_BALL) { + memcpy(q + adr, d->qpos + adr, (type == MJ_JNT_FREE ? 7 : 4)*sizeof(float)); + } else if (m->jnt_limited[j]) { + float lo = m->jnt_range[j][0]; + float hi = m->jnt_range[j][1]; + q[adr] = fminf(fmaxf((2.0f*d->qpos[adr] - lo - hi) / (hi - lo), -1.0f), 1.0f); + } else if (type == MJ_JNT_HINGE) { + q[adr] = fminf(fmaxf(d->qpos[adr] / (float)M_PI, -1.0f), 1.0f); + } else { + q[adr] = d->qpos[adr]; + } + } + memcpy(obs, q + skip, (m->nq - skip)*sizeof(float)); + obs += m->nq - skip; + for (int i = 0; i < m->nv; i++) { + obs[i] = fminf(fmaxf(d->qvel[i] / vel_scale, -1.0f), 1.0f); + } + return obs + m->nv; +} + +MJ_HD float* mj_obsScaled(float* obs, const float* src, int n, float scale) { + for (int i = 0; i < n; i++) { + obs[i] = fminf(fmaxf(src[i] / scale, -1.0f), 1.0f); + } + return obs + n; +} + // In-place Cholesky factor A = L L^T MJ_HD void mju_cholFactor(float* A, int n, int lda, int es) { for (int j = 0; j < n; j++) { @@ -499,15 +525,11 @@ MJ_HD void mj_local2Global(MjData* d, float* xpos, float* xmat, const float* pos } MJ_HD void mj_kinematics(const MjModel* m, MjData* d) { - memset(d->xpos[0], 0, 3*sizeof(float)); - memset(d->xmat[0], 0, 9*sizeof(float)); - d->xquat[0][0] = 1.0f; - d->xquat[0][1] = 0.0f; - d->xquat[0][2] = 0.0f; - d->xquat[0][3] = 0.0f; - d->xmat[0][0] = 1.0f; - d->xmat[0][4] = 1.0f; - d->xmat[0][8] = 1.0f; + // world body at the origin (ident also starts with the identity quaternion) + float ident[9] = {1.0f, 0.0f, 0.0f, 0.0f, 1.0f, 0.0f, 0.0f, 0.0f, 1.0f}; + memset(d->xpos[0], 0, sizeof(d->xpos[0])); + memcpy(d->xquat[0], ident, sizeof(d->xquat[0])); + memcpy(d->xmat[0], ident, sizeof(ident)); for (int i = 1; i < m->nbody; i++) { float xpos[3], xquat[4]; int jntadr = m->body_jntadr[i]; @@ -602,7 +624,8 @@ MJ_HD void mj_comPos(const MjModel* m, MjData* d) { memset(d->cinert[0], 0, 10*sizeof(float)); for (int i = 1; i < m->nbody; i++) { float* com = d->subtree_com[m->body_rootid[i]]; - float offset[3] = {d->xipos[i][0] - com[0], d->xipos[i][1] - com[1], d->xipos[i][2] - com[2]}; + float offset[3] = {d->xipos[i][0] - com[0], d->xipos[i][1] - com[1], + d->xipos[i][2] - com[2]}; mju_inertCom(d->cinert[i], m->body_inertia[i], d->ximat[i], offset, m->body_mass[i]); } for (int i = 1; i < m->nbody; i++) { @@ -745,9 +768,7 @@ MJ_HD void mj_passive(const MjModel* m, MjData* d) { continue; } int padr = m->jnt_qposadr[j]; - int dadr = m->jnt_dofadr[j]; - assert(m->jnt_type[j] >= MJ_JNT_SLIDE && "free/ball joint springs not supported"); - d->qfrc_passive[dadr] -= k*(d->qpos[padr] - m->qpos_spring[padr]); + d->qfrc_passive[m->jnt_dofadr[j]] -= k*(d->qpos[padr] - m->qpos_spring[padr]); } for (int i = 0; i < m->nv; i++) { d->qfrc_passive[i] -= m->dof_damping[i]*d->qvel[i]; @@ -944,17 +965,14 @@ MJ_HD void mj_collision(const MjModel* m, MjData* d) { } else if (t1 == MJ_GEOM_SPHERE && t2 == MJ_GEOM_SPHERE) { n = mjraw_SphereSphere(pre, margin, pos1, mat1, size1[0], pos2, mat2, size2[0]); } else if (t1 == MJ_GEOM_SPHERE && t2 == MJ_GEOM_CAPSULE) { - // sphere against the nearest point of the capsule segment float vec[3] = {pos1[0] - pos2[0], pos1[1] - pos2[1], pos1[2] - pos2[2]}; float x = fminf(fmaxf(mju_dot3(axis2, vec), -size2[1]), size2[1]); for (int k = 0; k < 3; k++) { vec[k] = pos2[k] + x*axis2[k]; } n = mjraw_SphereSphere(pre, margin, pos1, mat1, size1[0], vec, mat2, size2[0]); - } else if (t1 == MJ_GEOM_CAPSULE && t2 == MJ_GEOM_CAPSULE) { - n = mjc_CapsuleCapsule(pre, margin, pos1, mat1, size1, pos2, mat2, size2); } else { - assert(0 && "unsupported geom pair"); + n = mjc_CapsuleCapsule(pre, margin, pos1, mat1, size1, pos2, mat2, size2); } for (int c = 0; c < n && d->ncon < MJ_MAXCON; c++) { MjContact* con = &d->contact[d->ncon]; @@ -1027,22 +1045,16 @@ MJ_HD void mj_collision(const MjModel* m, MjData* d) { } } -// Constraints: rows of efc_J with pos/margin, MuJoCo's impedance and -// reference acceleration, then the dual QP for the constraint forces +// Constraints: efc rows, impedance/reference acceleration, dual QP solve -// Jacobian of a world point attached to body (jacp: translation, 3 x nv rows) +// Translational and rotational Jacobians (3 x nv each) of a world point on body MJ_HD void mj_jac(const MjModel* m, MjData* d, float jacp[3][MJ_MAX_NV], float jacr[3][MJ_MAX_NV], const float* point, int body) { memset(jacp, 0, 3*MJ_MAX_NV*sizeof(float)); - if (jacr) { - memset(jacr, 0, 3*MJ_MAX_NV*sizeof(float)); - } + memset(jacr, 0, 3*MJ_MAX_NV*sizeof(float)); float* com = d->subtree_com[m->body_rootid[body]]; float offset[3] = {point[0] - com[0], point[1] - com[1], point[2] - com[2]}; body = m->body_weldid[body]; - if (m->body_dofnum[body] == 0) { - return; - } for (int i = m->body_dofadr[body] + m->body_dofnum[body] - 1; i >= 0; i = m->dof_parentid[i]) { float* cdof = d->cdof[i]; @@ -1051,11 +1063,9 @@ MJ_HD void mj_jac(const MjModel* m, MjData* d, float jacp[3][MJ_MAX_NV], float j jacp[0][i] = cdof[3] + tmp[0]; jacp[1][i] = cdof[4] + tmp[1]; jacp[2][i] = cdof[5] + tmp[2]; - if (jacr) { - jacr[0][i] = cdof[0]; - jacr[1][i] = cdof[1]; - jacr[2][i] = cdof[2]; - } + jacr[0][i] = cdof[0]; + jacr[1][i] = cdof[1]; + jacr[2][i] = cdof[2]; } } @@ -1099,7 +1109,6 @@ MJ_HD void mj_makeConstraint(const MjModel* m, MjData* d) { if (!m->jnt_limited[j]) { continue; } - assert(m->jnt_type[j] == MJ_JNT_HINGE || m->jnt_type[j] == MJ_JNT_SLIDE); float value = d->qpos[m->jnt_qposadr[j]]; for (int side = -1; side <= 1; side += 2) { float dist = side*(m->jnt_range[j][(side + 1)/2] - value); @@ -1127,12 +1136,11 @@ MJ_HD void mj_makeConstraint(const MjModel* m, MjData* d) { break; } // Jacobian difference (body2 - body1) at the contact point, rotated - // into the contact frame: rows normal, tangent1, tangent2, and for - // condim > 3 the rotational rows torsion, roll1, roll2 + // into the contact frame; condim > 3 adds the rotational rows float jac1[3][MJ_MAX_NV], jac2[3][MJ_MAX_NV], jac[6][MJ_MAX_NV]; float jacr1[3][MJ_MAX_NV], jacr2[3][MJ_MAX_NV]; - mj_jac(m, d, jac1, dim > 3 ? jacr1 : NULL, con->pos, b1); - mj_jac(m, d, jac2, dim > 3 ? jacr2 : NULL, con->pos, b2); + mj_jac(m, d, jac1, jacr1, con->pos, b1); + mj_jac(m, d, jac2, jacr2, con->pos, b2); for (int r = 0; r < 3; r++) { for (int v = 0; v < m->nv; v++) { float dp = jac2[0][v] - jac1[0][v]; @@ -1186,19 +1194,14 @@ MJ_HD void mj_makeConstraint(const MjModel* m, MjData* d) { } } -// Constraint forces: min 1/2 f^T (A + R) f - f^T (aref - J qacc_smooth) with -// f >= 0 and A = J M^-1 J^T, by active-set (Lawson-Hanson NNLS) on the free -// set. Sets efc_force, qfrc_constraint and qacc. The dense matrices are the -// scratch bound by mj_makeData: efc_AR = A + R, efc_ARfree its factored -// free-set block and efc_MinvJT = M^-1 J^T, indexed with stride s. +// Constraint forces by active-set NNLS on the dual QP: min 1/2 f^T (A + R) f +// - f^T (aref - J qacc_smooth) with f >= 0 and A = J M^-1 J^T. efc_AR, +// efc_ARfree and efc_MinvJT are the strided scratch bound by mj_makeData. MJ_HD void mj_solveConstraint(const MjModel* m, MjData* d) { int nefc = d->nefc; int s = d->efc_stride; memcpy(d->qacc, d->qacc_smooth, sizeof(d->qacc)); memset(d->qfrc_constraint, 0, sizeof(d->qfrc_constraint)); - if (nefc == 0) { - return; - } float* MinvJT = d->efc_MinvJT; float* G = d->efc_AR; float* Gp = d->efc_ARfree; @@ -1298,9 +1301,8 @@ MJ_HD void mj_solveConstraint(const MjModel* m, MjData* d) { } } -// Body accelerations and interaction forces including constraint forces: -// cacc, cfrc_int and cfrc_ext (torque:force in the subtree COM frame; contact -// forces decoded from the pyramid rows) +// cacc, cfrc_int and cfrc_ext including constraint forces (torque:force in +// the subtree COM frame; contact forces decoded from the pyramid rows) MJ_HD void mj_rnePostConstraint(const MjModel* m, MjData* d) { memset(d->cfrc_ext, 0, sizeof(d->cfrc_ext)); for (int c = 0; c < d->ncon; c++) { @@ -1430,13 +1432,13 @@ MJ_HD void mj_RungeKutta4(const MjModel* m, MjData* d) { float B[4] = {1.0f/6.0f, 1.0f/3.0f, 1.0f/3.0f, 1.0f/6.0f}; float X[4][MJ_MAX_NQ + MJ_MAX_NV], F[4][MJ_MAX_NV]; float qpos0[MJ_MAX_NQ], time0 = d->time; + float dv[MJ_MAX_NV], da[MJ_MAX_NV]; memcpy(qpos0, d->qpos, sizeof(qpos0)); memcpy(X[0], d->qpos, sizeof(qpos0)); memcpy(X[0] + m->nq, d->qvel, sizeof(d->qvel)); memcpy(F[0], d->qacc, sizeof(d->qacc)); for (int i = 1; i < 4; i++) { // stage i uses only stage i-1 with weight A[i-1] - float dv[MJ_MAX_NV], da[MJ_MAX_NV]; for (int v = 0; v < m->nv; v++) { dv[v] = A[i - 1]*X[i - 1][m->nq + v]; da[v] = A[i - 1]*F[i - 1][v]; @@ -1451,7 +1453,6 @@ MJ_HD void mj_RungeKutta4(const MjModel* m, MjData* d) { memcpy(X[i] + m->nq, d->qvel, sizeof(d->qvel)); memcpy(F[i], d->qacc, sizeof(d->qacc)); } - float dv[MJ_MAX_NV], da[MJ_MAX_NV]; for (int v = 0; v < m->nv; v++) { dv[v] = 0.0f; da[v] = 0.0f; @@ -1473,7 +1474,6 @@ MJ_HD void mj_step(const MjModel* m, MjData* d) { if (m->opt_integrator == MJ_INT_RK4) { mj_RungeKutta4(m, d); } else { - assert(m->opt_integrator == MJ_INT_EULER); mj_Euler(m, d); } } diff --git a/ocean/mujoco/render.h b/ocean/mujoco/render.h index b0b974228e..72c566daaf 100644 --- a/ocean/mujoco/render.h +++ b/ocean/mujoco/render.h @@ -4,9 +4,9 @@ #include "raylib.h" #include "rlgl.h" -const Color MJ_BACKGROUND = (Color){6, 24, 24, 255}; -const Color MJ_BODY = (Color){0, 187, 187, 255}; -const Color MJ_CONTACT = (Color){187, 0, 0, 255}; +const Color MJ_BACKGROUND = {6, 24, 24, 255}; +const Color MJ_BODY = {0, 187, 187, 255}; +const Color MJ_CONTACT = {187, 0, 0, 255}; Vector3 mj_rl(const float* p) { return (Vector3){p[0], p[2], -p[1]}; @@ -39,8 +39,10 @@ void mj_render(const MjModel* m, MjData* d, const char* title, const float* targ float b[3] = {pos[0] + size[1]*mat[2], pos[1] + size[1]*mat[5], pos[2] + size[1]*mat[8]}; int type = m->geom_type[g]; if (type == MJ_GEOM_PLANE) { + // DrawGrid is finite: follow the camera in whole-tile steps rlPushMatrix(); - rlTranslatef(pos[0], pos[2], -pos[1]); + rlTranslatef(pos[0] + floorf(target[0] - pos[0]), pos[2], + -pos[1] - floorf(target[1] - pos[1])); DrawGrid(400, 1.0f); rlPopMatrix(); } else if (type == MJ_GEOM_SPHERE) { From 0a2773fda3a9403e7ea7cbac4e3379542fd3853b Mon Sep 17 00:00:00 2001 From: vyeoms Date: Mon, 7 Sep 2026 15:37:06 +0100 Subject: [PATCH 3/5] Add tests to check the MuJoCo implementation against python --- tests/mujoco_parity.py | 88 ++++++++++++++++++++ tests/test_mujoco_gpu.cu | 111 +++++++++++++++++++++++++ tests/test_mujoco_parity.c | 161 +++++++++++++++++++++++++++++++++++++ 3 files changed, 360 insertions(+) create mode 100644 tests/mujoco_parity.py create mode 100644 tests/test_mujoco_gpu.cu create mode 100644 tests/test_mujoco_parity.c diff --git a/tests/mujoco_parity.py b/tests/mujoco_parity.py new file mode 100644 index 0000000000..b4d08b3a76 --- /dev/null +++ b/tests/mujoco_parity.py @@ -0,0 +1,88 @@ +"""Parity test: ocean/mujoco envs vs MuJoCo through gymnasium. + +Writes a reference trajectory (random actions) with gymnasium + mujoco, compiles +tests/test_mujoco_parity.c against ocean/mujoco/mjc_.h and runs it with +resources/mujoco/.bin. Requires pip install "gymnasium[mujoco]". + + python tests/mujoco_parity.py mjc_half_cheetah [--steps 300] [--seed 123] +""" +import argparse +import os +import subprocess +import sys +import tempfile + +import numpy as np + +GYM_IDS = {"half_cheetah": "HalfCheetah-v5", "hopper": "Hopper-v5", "walker2d": "Walker2d-v5", + "ant": "Ant-v5", "humanoid": "Humanoid-v5", "swimmer": "Swimmer-v5"} +# reward weights passed to the harness (gymnasium defaults) +DEFINES = {"half_cheetah": ["-DCTRL_COST_WEIGHT=0.1"], + "hopper": ["-DCTRL_COST_WEIGHT=1e-3", "-DHEALTHY_REWARD=1.0"], + "walker2d": ["-DCTRL_COST_WEIGHT=1e-3", "-DHEALTHY_REWARD=1.0"], + "ant": ["-DCTRL_COST_WEIGHT=0.5", "-DHEALTHY_REWARD=1.0", "-DCONTACT_COST_WEIGHT=5e-4"], + "humanoid": ["-DCTRL_COST_WEIGHT=0.1", "-DHEALTHY_REWARD=5.0", "-DCONTACT_COST_WEIGHT=5e-7", + "-DFORWARD_REWARD_WEIGHT=1.25"], + "swimmer": ["-DCTRL_COST_WEIGHT=1e-4"]} + + +def write_reference(path, gym_id, steps, seed): + import gymnasium as gym + env = gym.make(gym_id) + env.reset(seed=seed) + u = env.unwrapped + m, d = u.model, u.data + rng = np.random.default_rng(seed) + with open(path, "w") as f: + f.write("%d %d %d %d %d\n" % (steps, m.nq, m.nv, m.nu, u.observation_space.shape[0])) + for _ in range(steps): + # start state, action, resulting state, reward, ncon, terminated, obs + f.write(" ".join("%.17g" % v for v in d.qpos) + "\n") + f.write(" ".join("%.17g" % v for v in d.qvel) + "\n") + a = rng.uniform(u.action_space.low, u.action_space.high) + obs, r, term, trunc, _ = env.step(a) + f.write(" ".join("%.17g" % v for v in a) + "\n") + f.write(" ".join("%.17g" % v for v in d.qpos) + "\n") + f.write(" ".join("%.17g" % v for v in d.qvel) + "\n") + f.write("%.17g %d %d\n" % (r, d.ncon, term)) + f.write(" ".join("%.17g" % v for v in obs) + "\n") + if term: + env.reset(seed=int(rng.integers(1 << 30))) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("env") + ap.add_argument("--steps", type=int, default=300) + ap.add_argument("--seed", type=int, default=123) + ap.add_argument("--ref", help="write/keep the reference file here") + ap.add_argument("--cc", default=os.environ.get("CC", "clang")) + args = ap.parse_args() + name = args.env[4:] if args.env.startswith("mjc_") else args.env + root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + raylib = os.path.join(root, "raylib-5.5_linux_amd64") + with tempfile.TemporaryDirectory() as tmp: + ref = args.ref or os.path.join(tmp, "ref.txt") + exe = os.path.join(tmp, "test_mujoco_parity") + write_reference(ref, GYM_IDS[name], args.steps, args.seed) + subprocess.check_call([ + args.cc, "-O2", "-Wno-narrowing", "-Wno-unused-function", + "-DENV_HEADER=\"../ocean/mujoco/mjc_%s.h\"" % name] + DEFINES[name] + [ + "-I" + os.path.join(root, "src"), "-I" + os.path.join(root, "ocean", "mujoco"), + "-I" + os.path.join(raylib, "include"), + os.path.join(root, "tests", "test_mujoco_parity.c"), + os.path.join(raylib, "lib", "libraylib.a"), + "-lGL", "-lm", "-lpthread", "-ldl", "-o", exe]) + model = os.path.join(root, "resources", "mujoco", name + ".bin") + out = subprocess.check_output([exe, ref, model], text=True, cwd=root) + print(out) + first = out.splitlines()[0] + max_dq = float(first.split("max|dq|")[1].split()[0]) + max_dv = float(first.split("max|dv|")[1].split()[0]) + ok = max_dq < 1e-4 and max_dv < 1e-2 + print("PARITY", "PASS" if ok else "FAIL") + sys.exit(0 if ok else 1) + + +if __name__ == "__main__": + main() diff --git a/tests/test_mujoco_gpu.cu b/tests/test_mujoco_gpu.cu new file mode 100644 index 0000000000..e9684bc41b --- /dev/null +++ b/tests/test_mujoco_gpu.cu @@ -0,0 +1,111 @@ +// CPU vs GPU parity of an ocean/mujoco env: the same envs (same seeds and +// actions) are stepped by mjc_step on the host and by the --cu backend kernels, +// and qpos/qvel/obs/reward are compared after every step. Build with +// -DENV_HEADER='"../ocean/mujoco/mjc_ENV.h"'. Usage: test_mujoco_gpu MODEL_BIN [N] [STEPS] +#include +#include +#define PUF_BACKEND PUF_GPU +typedef float obs_t; +#include ENV_HEADER + +int main(int argc, char** argv) { + int n = argc > 2 ? atoi(argv[2]) : 256; + int steps = argc > 3 ? atoi(argv[3]) : 200; + Dict kwargs = {0}; + dict_set_str(&kwargs, "model", argv[1]); + dict_set(&kwargs, "max_steps", 50); + dict_set(&kwargs, "reset_noise_scale", 0.1); + dict_set(&kwargs, "forward_reward_weight", 1.0); + dict_set(&kwargs, "ctrl_cost_weight", 0.1); + dict_set(&kwargs, "contact_cost_weight", 5e-4); + dict_set(&kwargs, "healthy_reward", 1.0); + obs_t* d_obs; + float* d_act; + float* d_rew; + float* d_term; + cudaMalloc((void**)&d_obs, n*OBS_SIZE*sizeof(obs_t)); + cudaMalloc((void**)&d_act, n*NUM_ATNS*sizeof(float)); + cudaMalloc((void**)&d_rew, n*sizeof(float)); + cudaMalloc((void**)&d_term, n*sizeof(float)); + Env* d_envs = puf_vec_create(n, &kwargs, d_obs, d_act, d_rew, d_term); + // host twins with the same seeds + Env* envs = (Env*)calloc(n, sizeof(Env)); + float* obs = (float*)calloc(n*OBS_SIZE, sizeof(float)); + float* act = (float*)calloc(n*NUM_ATNS, sizeof(float)); + float* rew = (float*)calloc(n, sizeof(float)); + float* term = (float*)calloc(n, sizeof(float)); + for (int i = 0; i < n; i++) { + mjc_init(&envs[i], &kwargs); + mj_makeData(&envs[i].d, (float*)calloc(MJ_SCRATCH, sizeof(float)), 1); + envs[i].rng = i + 1; + envs[i].agents[0].observations = obs + i*OBS_SIZE; + envs[i].agents[0].actions = act + i*NUM_ATNS; + envs[i].agents[0].rewards = rew + i; + envs[i].agents[0].terminals = term + i; + mjc_reset(&envs[i]); + } + puf_reset(d_envs); + cudaDeviceSynchronize(); + size_t free0, total; + cudaMemGetInfo(&free0, &total); + float* g_obs = (float*)calloc(n*OBS_SIZE, sizeof(float)); + float* g_rew = (float*)calloc(n, sizeof(float)); + float* g_term = (float*)calloc(n, sizeof(float)); + Env* g_env = (Env*)calloc(1, sizeof(Env)); + double max_o = 0, max_r = 0, max_q = 0, sum_r = 0; + int term_mismatch = 0; + double gpu_ms = 0; + for (int t = 0; t < steps; t++) { + for (int i = 0; i < n*NUM_ATNS; i++) { + act[i] = 2.0f*((float)rand() / RAND_MAX) - 1.0f; + } + cudaMemcpy(d_act, act, n*NUM_ATNS*sizeof(float), cudaMemcpyHostToDevice); + cudaEvent_t e0, e1; + cudaEventCreate(&e0); + cudaEventCreate(&e1); + cudaEventRecord(e0); + puf_step(d_envs); + cudaEventRecord(e1); + cudaDeviceSynchronize(); + float ms; + cudaEventElapsedTime(&ms, e0, e1); + gpu_ms += ms; + for (int i = 0; i < n; i++) { + mjc_step(&envs[i]); + } + cudaMemcpy(g_obs, d_obs, n*OBS_SIZE*sizeof(float), cudaMemcpyDeviceToHost); + cudaMemcpy(g_rew, d_rew, n*sizeof(float), cudaMemcpyDeviceToHost); + cudaMemcpy(g_term, d_term, n*sizeof(float), cudaMemcpyDeviceToHost); + cudaMemcpy(g_env, d_envs, sizeof(Env), cudaMemcpyDeviceToHost); + for (int i = 0; i < n*OBS_SIZE; i++) { + max_o = fmax(max_o, fabs(g_obs[i] - obs[i])); + } + for (int i = 0; i < n; i++) { + max_r = fmax(max_r, fabs(g_rew[i] - rew[i])); + sum_r += fabs(g_rew[i] - rew[i]); + term_mismatch += (g_term[i] != 0) != (term[i] != 0); + } + double eq = 0, eo = 0; + for (int i = 0; i < mj_model.nq; i++) { + eq = fmax(eq, fabs(g_env->d.qpos[i] - envs[0].d.qpos[i])); + } + for (int i = 0; i < n*OBS_SIZE; i++) { + eo = fmax(eo, fabs(g_obs[i] - obs[i])); + } + max_q = fmax(max_q, eq); + int k = t + 1; + if (k == 1 || k == 2 || k == 5 || k == 10 || k == 20 || k == 50 || k == steps) { + printf("step %4d: max|dq|(env0) %.3e max|dobs|(all) %.3e\n", k, eq, eo); + } + } + printf("cpu-vs-gpu over %d envs x %d steps: max|dobs| %.3e max|dr| %.3e mean|dr| %.3e " + "max|dq|(env0) %.3e terminal mismatch %d\n", n, steps, max_o, max_r, sum_r / (n*steps), + max_q, term_mismatch); + printf("gpu step kernel: %.2f ms per batch of %d (%.0f env steps/s)\n", gpu_ms / steps, n, + n*steps / (gpu_ms / 1000.0)); + size_t free1; + cudaMemGetInfo(&free1, &total); + printf("device memory taken by the step kernel launch (local memory): %.0f MB; envs %.0f MB\n", + (free0 - free1) / 1048576.0, n*sizeof(Env) / 1048576.0); + return 0; +} diff --git a/tests/test_mujoco_parity.c b/tests/test_mujoco_parity.c new file mode 100644 index 0000000000..21120d1b44 --- /dev/null +++ b/tests/test_mujoco_parity.c @@ -0,0 +1,161 @@ +// Parity check of an ocean/mujoco env against a MuJoCo (gymnasium) reference +// trajectory written by tests/mujoco_parity.py. Build with +// -DENV_HEADER='"../ocean/mujoco/mjc_ENV.h"' and the env's reward weights +// (-DCTRL_COST_WEIGHT=..., -DHEALTHY_REWARD=..., -DCONTACT_COST_WEIGHT=..., +// -DFORWARD_REWARD_WEIGHT=...). +// Usage: test_mujoco_parity REF_FILE MODEL_BIN +#include +#include +#include ENV_HEADER + +#define MAXT 4096 +#ifndef FORWARD_REWARD_WEIGHT +#define FORWARD_REWARD_WEIGHT 1.0f +#endif + +double ref_q0[MAXT][MJ_MAX_NQ], ref_v0[MAXT][MJ_MAX_NV], ref_a[MAXT][MJ_MAX_NU]; +double ref_q[MAXT][MJ_MAX_NQ], ref_v[MAXT][MJ_MAX_NV]; +double ref_r[MAXT], ref_obs[MAXT][OBS_SIZE]; +int ref_ncon[MAXT], ref_term[MAXT]; + +void set_state(Env* env, int t) { + for (int i = 0; i < mj_model.nq; i++) env->d.qpos[i] = ref_q0[t][i]; + for (int i = 0; i < mj_model.nv; i++) env->d.qvel[i] = ref_v0[t][i]; + mj_kinematics(&mj_model, &env->d); + env->tick = 0; +} + +int main(int argc, char** argv) { + mj_loadModel(&mj_model, argv[2]); + FILE* fp = fopen(argv[1], "r"); + int T, nq, nv, nu, nobs; + fscanf(fp, "%d %d %d %d %d", &T, &nq, &nv, &nu, &nobs); + if (nq != mj_model.nq || nv != mj_model.nv || nu != mj_model.nu || nobs != OBS_SIZE) { + printf("size mismatch: ref nq %d nv %d nu %d nobs %d vs model %d %d %d %d\n", + nq, nv, nu, nobs, mj_model.nq, mj_model.nv, mj_model.nu, OBS_SIZE); + return 1; + } + for (int t = 0; t < T; t++) { + for (int i = 0; i < mj_model.nq; i++) fscanf(fp, "%lf", &ref_q0[t][i]); + for (int i = 0; i < mj_model.nv; i++) fscanf(fp, "%lf", &ref_v0[t][i]); + for (int i = 0; i < mj_model.nu; i++) fscanf(fp, "%lf", &ref_a[t][i]); + for (int i = 0; i < mj_model.nq; i++) fscanf(fp, "%lf", &ref_q[t][i]); + for (int i = 0; i < mj_model.nv; i++) fscanf(fp, "%lf", &ref_v[t][i]); + fscanf(fp, "%lf %d %d", &ref_r[t], &ref_ncon[t], &ref_term[t]); + for (int i = 0; i < OBS_SIZE; i++) fscanf(fp, "%lf", &ref_obs[t][i]); + } + fclose(fp); + + float obs[OBS_SIZE], actions[MJ_MAX_NU], reward, terminal; + Env env = {0}; + env.m = &mj_model; + mj_makeData(&env.d, (float*)calloc(MJ_SCRATCH, sizeof(float)), 1); + env.agents[0].observations = obs; + env.agents[0].actions = actions; + env.agents[0].rewards = &reward; + env.agents[0].terminals = &terminal; + env.num_agents = 1; + env.max_steps = 1 << 30; + env.reset_noise_scale = 0.1f; + env.forward_reward_weight = FORWARD_REWARD_WEIGHT; + env.ctrl_cost_weight = CTRL_COST_WEIGHT; +#ifdef HEALTHY_REWARD + env.healthy_reward = HEALTHY_REWARD; +#endif +#ifdef CONTACT_COST_WEIGHT + env.contact_cost_weight = CONTACT_COST_WEIGHT; +#endif + Env env0 = env; + puf_reset(&env); + + // One-step parity: start every step from the reference start state. The env + // resets itself on termination, so terminal steps only compare the flag. + double max_q = 0, max_v = 0, max_r = 0, sum_q = 0, sum_v = 0; + int arg_q = -1, arg_v = -1, dof_q = -1, dof_v = -1, ncon_mismatch = 0, term_mismatch = 0; + int arg_r = -1; + int ncmp = 0; + for (int t = 0; t < T; t++) { + set_state(&env, t); + for (int i = 0; i < mj_model.nu; i++) actions[i] = ref_a[t][i]; + double ret0 = env.episode_return; + puf_step(&env); + // the trainer reward is scaled; episode_return accumulates the Gym reward + double gym_reward = env.episode_return - ret0; + term_mismatch += (terminal != 0) != (ref_term[t] != 0); + if (ref_term[t] || terminal) { + continue; + } + ncmp++; + for (int i = 0; i < mj_model.nq; i++) { + double eq = fabs(env.d.qpos[i] - ref_q[t][i]); + sum_q += eq; + if (eq > max_q) { max_q = eq; arg_q = t; dof_q = i; } + } + for (int i = 0; i < mj_model.nv; i++) { + double ev = fabs(env.d.qvel[i] - ref_v[t][i]); + sum_v += ev; + if (ev > max_v) { max_v = ev; arg_v = t; dof_v = i; } + } + double er = fabs(gym_reward - ref_r[t]); + if (er > max_r) { max_r = er; arg_r = t; } + if (env.d.ncon != ref_ncon[t]) ncon_mismatch++; + } + printf("one-step: max|dq| %.3e (t=%d i=%d) max|dv| %.3e (t=%d i=%d)\n", + max_q, arg_q, dof_q, max_v, arg_v, dof_v); + printf("one-step: mean|dq| %.3e mean|dv| %.3e max|dr| %.3e (t=%d) ncon mismatch" + " %d/%d terminal mismatch %d\n", sum_q / (ncmp*mj_model.nq), sum_v / (ncmp*mj_model.nv), + max_r, arg_r, ncon_mismatch, ncmp, term_mismatch); + + // Free run from the initial state until the reference episode ends. Reward + // is compared while the trajectories are still in sync. + set_state(&env, 0); + double ret = 0, ref_ret = 0, sync_dr = 0; + int nsync = 0; + printf("free-run t | max|dq| | max|dv| | qpos[0] vs ref\n"); + for (int t = 0; t < T; t++) { + if (ref_term[t]) { + printf("reference episode terminated at step %d, stopping free run\n", t + 1); + break; + } + for (int i = 0; i < mj_model.nu; i++) actions[i] = ref_a[t][i]; + double ret0 = env.episode_return; + puf_step(&env); + double gym_reward = env.episode_return - ret0; + ret += gym_reward; + ref_ret += ref_r[t]; + double eq = 0, ev = 0; + for (int i = 0; i < mj_model.nq; i++) eq = fmax(eq, fabs(env.d.qpos[i] - ref_q[t][i])); + for (int i = 0; i < mj_model.nv; i++) ev = fmax(ev, fabs(env.d.qvel[i] - ref_v[t][i])); + if (eq < 1e-4) { + sync_dr = fmax(sync_dr, fabs(gym_reward - ref_r[t])); + nsync++; + } + int n = t + 1; + if (n == 1 || n == 2 || n == 5 || n == 10 || n == 20 || n == 50 || n == 100 || n == 200 + || n == T) { + printf("%9d | %.3e | %.3e | %.4f vs %.4f\n", n, eq, ev, env.d.qpos[0], ref_q[t][0]); + } + } + printf("free-run return %.3f vs ref %.3f; max|dr| %.3e over %d in-sync steps\n", ret, ref_ret, + sync_dr, nsync); + + // Throughput with random actions + env = env0; + env.max_steps = 1000; + puf_reset(&env); + struct timespec t0, t1; + clock_gettime(CLOCK_MONOTONIC, &t0); + int N = 20000; + int bad = 0; + for (int t = 0; t < N; t++) { + for (int i = 0; i < mj_model.nu; i++) actions[i] = 2.0f*((float)rand() / RAND_MAX) - 1.0f; + puf_step(&env); + bad += !isfinite(env.d.qpos[0]) || !isfinite(reward); + } + clock_gettime(CLOCK_MONOTONIC, &t1); + double dt = (t1.tv_sec - t0.tv_sec) + 1e-9*(t1.tv_nsec - t0.tv_nsec); + printf("random policy: %.0f steps/s (single thread), %d episodes, mean return %.1f, " + "nonfinite %d\n", N / dt, (int)env.log.n, env.log.n > 0 ? env.log.score / env.log.n : 0.0f, + bad); + return 0; +} From a067213f2104384ba18c2c774ed925c1dca2ac78 Mon Sep 17 00:00:00 2001 From: vyeoms Date: Wed, 9 Sep 2026 12:14:20 +0100 Subject: [PATCH 4/5] Renormalize Humanoid range to -1,1 --- ocean/mujoco/mjc_humanoid.h | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/ocean/mujoco/mjc_humanoid.h b/ocean/mujoco/mjc_humanoid.h index 6fa6facaf7..9365f1b064 100644 --- a/ocean/mujoco/mjc_humanoid.h +++ b/ocean/mujoco/mjc_humanoid.h @@ -117,8 +117,9 @@ MJ_HD void mjc_step(Humanoid* env) { float* actions = env->agents[0].actions; float cost = 0.0f; for (int i = 0; i < m->nu; i++) { + // policy actions rescaled to +-1 (gym clips at +-0.4) const float* range = m->actuator_ctrlrange[i]; - env->d.ctrl[i] = fminf(fmaxf(actions[i], range[0]), range[1]); + env->d.ctrl[i] = range[1]*fminf(fmaxf(actions[i], -1.0f), 1.0f); cost += env->d.ctrl[i]*env->d.ctrl[i]; } float x0 = mjc_com_x(env); From 268d731bec3094db87bce67e7a9e42d6a620efda Mon Sep 17 00:00:00 2001 From: vyeoms Date: Wed, 9 Sep 2026 16:10:12 +0100 Subject: [PATCH 5/5] Fix eval framerate and clean comments --- ocean/mujoco/backend.h | 39 ++++++++++++++++++++++++++++----- ocean/mujoco/mjc_ant.h | 1 - ocean/mujoco/mjc_half_cheetah.h | 1 - ocean/mujoco/mjc_hopper.h | 1 - ocean/mujoco/mjc_humanoid.h | 1 - ocean/mujoco/mjc_walker2d.h | 1 - 6 files changed, 34 insertions(+), 10 deletions(-) diff --git a/ocean/mujoco/backend.h b/ocean/mujoco/backend.h index e7c28663bf..f8134fd61c 100644 --- a/ocean/mujoco/backend.h +++ b/ocean/mujoco/backend.h @@ -1,5 +1,34 @@ // Vector backend, included at the end of each mjc_.h. +#ifndef PUFFERCPU_EVAL_MAIN + +static void mjc_render_paced(Env* env) { + int nq = env->m->nq; + static float prev[MJ_MAX_NQ]; + static int have_prev = 0; + static float credit = 0.0f; + credit += MJC_FRAME_SKIP*env->m->opt_timestep*60.0f; + int frames = (int)credit; + credit -= frames; + if (have_prev && frames > 1) { + float cur[MJ_MAX_NQ]; + memcpy(cur, env->d.qpos, nq*sizeof(float)); + for (int f = 1; f <= frames; f++) { + float a = (float)f/frames; + for (int i = 0; i < nq; i++) { + env->d.qpos[i] = prev[i] + a*(cur[i] - prev[i]); + } + mjc_render(env); + } + memcpy(env->d.qpos, cur, nq*sizeof(float)); + } else if (!have_prev || frames >= 1) { + mjc_render(env); + } + memcpy(prev, env->d.qpos, nq*sizeof(float)); + have_prev = 1; +} +#endif + #if PUF_BACKEND == PUF_GPU #define MJC_BLOCK 128 #define MJC_STATE (offsetof(Env, d) + offsetof(MjData, xquat)) @@ -80,7 +109,7 @@ void puf_render(Env* envs) { cudaStreamSynchronize(mjc_gpu.stream); cudaMemcpy(&env, mjc_gpu.envs, sizeof(Env), cudaMemcpyDeviceToHost); env.m = &mj_model; - mjc_render(&env); + mjc_render_paced(&env); } void puf_close(Env* envs) { @@ -96,10 +125,6 @@ void puf_init(Env* env, Dict* kwargs) { } #ifdef PUFFERCPU_EVAL_MAIN -// Whole control steps (MJC_FRAME_SKIP substeps at once) strobe on screen at -// 20 Hz. Run the real mjc_step for exact rewards, logs and resets, then -// rewind and re-integrate the same deterministic substeps one per rendered -// frame; PUF_STEPS_PER_SEC is therefore the substep rate 1/opt_timestep. #define PUF_EVAL_SHOULD_FORWARD #define MJC_STATE (offsetof(MjData, xquat)) char mjc_true[MJC_STATE]; @@ -139,7 +164,11 @@ void puf_step(Env* env) { #endif void puf_render(Env* env) { +#ifdef PUFFERCPU_EVAL_MAIN mjc_render(env); +#else + mjc_render_paced(env); +#endif } void puf_close(Env* env) { diff --git a/ocean/mujoco/mjc_ant.h b/ocean/mujoco/mjc_ant.h index fd9380bd0d..d38d582802 100644 --- a/ocean/mujoco/mjc_ant.h +++ b/ocean/mujoco/mjc_ant.h @@ -9,7 +9,6 @@ #include "raylib.h" typedef float obs_t; #include "pufferenv.h" -// model capacities, checked by mj_loadModel #define MJ_MAX_NQ 15 #define MJ_MAX_NV 14 #define MJ_MAX_NBODY 14 diff --git a/ocean/mujoco/mjc_half_cheetah.h b/ocean/mujoco/mjc_half_cheetah.h index 95c9973b1b..bebd7b936d 100644 --- a/ocean/mujoco/mjc_half_cheetah.h +++ b/ocean/mujoco/mjc_half_cheetah.h @@ -8,7 +8,6 @@ #include "raylib.h" typedef float obs_t; #include "pufferenv.h" -// model capacities, checked by mj_loadModel #define MJ_MAX_NQ 9 #define MJ_MAX_NV 9 #define MJ_MAX_NBODY 8 diff --git a/ocean/mujoco/mjc_hopper.h b/ocean/mujoco/mjc_hopper.h index 594333610c..83a8a6eaf5 100644 --- a/ocean/mujoco/mjc_hopper.h +++ b/ocean/mujoco/mjc_hopper.h @@ -8,7 +8,6 @@ #include "raylib.h" typedef float obs_t; #include "pufferenv.h" -// model capacities, checked by mj_loadModel #define MJ_MAX_NQ 6 #define MJ_MAX_NV 6 #define MJ_MAX_NBODY 5 diff --git a/ocean/mujoco/mjc_humanoid.h b/ocean/mujoco/mjc_humanoid.h index 9365f1b064..aab45b13f0 100644 --- a/ocean/mujoco/mjc_humanoid.h +++ b/ocean/mujoco/mjc_humanoid.h @@ -10,7 +10,6 @@ #include "raylib.h" typedef float obs_t; #include "pufferenv.h" -// model capacities, checked by mj_loadModel #define MJ_MAX_NQ 24 #define MJ_MAX_NV 23 #define MJ_MAX_NBODY 14 diff --git a/ocean/mujoco/mjc_walker2d.h b/ocean/mujoco/mjc_walker2d.h index 57f98426bc..5134c2b2b7 100644 --- a/ocean/mujoco/mjc_walker2d.h +++ b/ocean/mujoco/mjc_walker2d.h @@ -9,7 +9,6 @@ #include "raylib.h" typedef float obs_t; #include "pufferenv.h" -// model capacities, checked by mj_loadModel #define MJ_MAX_NQ 9 #define MJ_MAX_NV 9 #define MJ_MAX_NBODY 8