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..7bcba27dbd --- /dev/null +++ b/config/mjc_ant.ini @@ -0,0 +1,107 @@ +[base] +env_name = mjc_ant + +[vec] +total_agents = 256 +num_buffers = 1 +num_threads = 16 + +[env] +model = resources/mujoco/ant.bin +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 = 512 +num_layers = 2 +expansion_factor = 1 + +[train] +gpus = 1 +total_timesteps = 10000000 +learning_rate = 0.003 +anneal_lr = 1 +min_lr_ratio = 0 +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 = 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 new file mode 100644 index 0000000000..9b6b47ccc9 --- /dev/null +++ b/config/mjc_half_cheetah.ini @@ -0,0 +1,105 @@ +[base] +env_name = mjc_half_cheetah + +[vec] +total_agents = 512 +num_buffers = 1 +num_threads = 16 + +[env] +model = resources/mujoco/half_cheetah.bin +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 +total_timesteps = 10000000 +learning_rate = 0.0036 +anneal_lr = 1 +min_lr_ratio = 0 +gamma = 0.975 +gae_lambda = 0.922 +replay_ratio = 4.9 +clip_coef = 0.11 +vf_coef = 0.5 +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 new file mode 100644 index 0000000000..634528ca9c --- /dev/null +++ b/config/mjc_hopper.ini @@ -0,0 +1,106 @@ +[base] +env_name = mjc_hopper + +[vec] +total_agents = 256 +num_buffers = 1 +num_threads = 16 + +[env] +model = resources/mujoco/hopper.bin +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 +total_timesteps = 10000000 +learning_rate = 0.0017 +anneal_lr = 1 +min_lr_ratio = 0 +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 new file mode 100644 index 0000000000..9a0a94d34b --- /dev/null +++ b/config/mjc_humanoid.ini @@ -0,0 +1,107 @@ +[base] +env_name = mjc_humanoid + +[vec] +total_agents = 512 +num_buffers = 1 +num_threads = 16 + +[env] +model = resources/mujoco/humanoid.bin +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 +total_timesteps = 10000000 +learning_rate = 0.00093 +anneal_lr = 1 +min_lr_ratio = 0 +gamma = 0.99 +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.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 new file mode 100644 index 0000000000..e72f6d626d --- /dev/null +++ b/config/mjc_walker2d.ini @@ -0,0 +1,106 @@ +[base] +env_name = mjc_walker2d + +[vec] +total_agents = 256 +num_buffers = 1 +num_threads = 16 + +[env] +model = resources/mujoco/walker2d.bin +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 +total_timesteps = 10000000 +learning_rate = 0.004 +anneal_lr = 1 +min_lr_ratio = 0 +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 new file mode 100644 index 0000000000..f8134fd61c --- /dev/null +++ b/ocean/mujoco/backend.h @@ -0,0 +1,179 @@ +// 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)) + +struct { + Env* envs; + int n; + cudaStream_t stream; +} mjc_gpu; + +__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); + if (reset) { + mjc_reset(&env); + } else { + 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, + float* terminals) { + Env* host = (Env*)calloc(n, sizeof(Env)); + MjModel* m; + assert(cudaMalloc((void**)&m, sizeof(MjModel)) == cudaSuccess); + float* scratch; + 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]; + 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_launch(1); +} + +void puf_step(Env* envs) { + mjc_launch(0); +} + +// Draw env 0 with the host model (mj_render recomputes kinematics 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_paced(&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); +} + +#ifdef PUFFERCPU_EVAL_MAIN +#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); +} + +void puf_step(Env* env) { + mjc_step(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) { + 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..d38d582802 --- /dev/null +++ b/ocean/mujoco/mjc_ant.h @@ -0,0 +1,171 @@ +// 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 +#include +#include +#include "raylib.h" +typedef float obs_t; +#include "pufferenv.h" +#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" + +// 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 100 + +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; + int tick_frames_left; + 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* 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++) { + 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 displacement with body xpos, which lags qpos by one substep + float x0 = env->d.xpos[1][0]; + 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 = 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); + 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 / ANT_TARGET_VEL; + 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 / 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; + 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->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..bebd7b936d --- /dev/null +++ b/ocean/mujoco/mjc_half_cheetah.h @@ -0,0 +1,150 @@ +// HalfCheetah (gymnasium HalfCheetah-v5): obs = qpos[1:] + qvel scaled by +// mj_obsJoints, reward = forward velocity - ctrl cost, 1000 step episodes. + +#include +#include +#include +#include +#include "raylib.h" +typedef float obs_t; +#include "pufferenv.h" +#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" + +// 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 100 + +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; + int tick_frames_left; + 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 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; + mj_obsJoints(m, &env->d, env->agents[0].observations, 1, 20.0f); +} + +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 < MJC_FRAME_SKIP; k++) { + mj_step(m, &env->d); + } + 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 / HC_TARGET_VEL; + env->agents[0].terminals[0] = 0.0f; + if (env->tick < env->max_steps) { + 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 / 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; + 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->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..83a8a6eaf5 --- /dev/null +++ b/ocean/mujoco/mjc_hopper.h @@ -0,0 +1,155 @@ +// Hopper (gymnasium Hopper-v5): obs = qpos[1:] + qvel scaled by mj_obsJoints, +// reward = healthy + forward velocity - ctrl cost, terminates when unhealthy. + +#include +#include +#include +#include +#include "raylib.h" +typedef float obs_t; +#include "pufferenv.h" +#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" + +// 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 500 + +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; + int tick_frames_left; + 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 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; + mj_obsJoints(m, &env->d, env->agents[0].observations, 1, 10.0f); +} + +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 < MJC_FRAME_SKIP; k++) { + mj_step(m, &env->d); + } + float* q = env->d.qpos; + 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++) { + 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 / HP_TARGET_VEL; + env->agents[0].terminals[0] = 0.0f; + if (healthy && env->tick < env->max_steps) { + 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 / 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; + 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->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..aab45b13f0 --- /dev/null +++ b/ocean/mujoco/mjc_humanoid.h @@ -0,0 +1,189 @@ +// 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 +#include +#include +#include "raylib.h" +typedef float obs_t; +#include "pufferenv.h" +#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" + +// 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 333 + +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; + int tick_frames_left; + 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 = 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]); + } + return fminf(env->contact_cost_weight*cost, HM_CONTACT_COST_MAX); +} + +// 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; + 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++) { + // policy actions rescaled to +-1 (gym clips at +-0.4) + const float* range = m->actuator_ctrlrange[i]; + 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); + 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 = 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); + 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 / HM_TARGET_VEL; + 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 / 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; + 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->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..5134c2b2b7 --- /dev/null +++ b/ocean/mujoco/mjc_walker2d.h @@ -0,0 +1,150 @@ +// 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 +#include +#include +#include "raylib.h" +typedef float obs_t; +#include "pufferenv.h" +#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" + +// 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 500 + +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; + int tick_frames_left; + 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 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; + mj_obsJoints(m, &env->d, env->agents[0].observations, 1, 10.0f); +} + +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 < MJC_FRAME_SKIP; k++) { + mj_step(m, &env->d); + } + float* q = env->d.qpos; + 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 / WK_TARGET_VEL; + env->agents[0].terminals[0] = 0.0f; + if (healthy && env->tick < env->max_steps) { + 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 / 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; + 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->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..d659538112 --- /dev/null +++ b/ocean/mujoco/mjcf2bin.py @@ -0,0 +1,77 @@ +"""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" + 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" + 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..e41ae9ded1 --- /dev/null +++ b/ocean/mujoco/physics.h @@ -0,0 +1,1484 @@ +// MuJoCo-style rigid body physics for PufferLib envs. + +#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 // 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}; +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; + +#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[2]; + mj_read(fp, head, 2, sizeof(int)); + assert(head[0] == MJ_MAGIC && head[1] == 1 && "bad model file"); + 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"); + 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; + } +} + +// res may alias qa (kinematics accumulates joint rotations in place) +MJ_HD void mju_mulQuat(float* res, const float* qa, const float* qb) { + 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) { + 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); +} + +// 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++) { + 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) { + // 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]; + 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]; + 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]; + } +} + +// 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) { + 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 { + 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]; + 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: efc rows, impedance/reference acceleration, dual QP solve + +// 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)); + 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]; + 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]; + 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; + } + 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; 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, 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]; + 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 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)); + 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]; + } + } +} + +// 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++) { + 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; + 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] + 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)); + } + 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 { + 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..72c566daaf --- /dev/null +++ b/ocean/mujoco/render.h @@ -0,0 +1,71 @@ +// 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 = {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]}; +} + +// 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) { + // DrawGrid is finite: follow the camera in whole-tile steps + rlPushMatrix(); + 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) { + 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 @@ + + + + + + + + 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; +}