From 77ccfd42686855b5b373e22d6ed2a74829420d35 Mon Sep 17 00:00:00 2001 From: Zilong Date: Sat, 15 Aug 2026 14:25:29 -0500 Subject: [PATCH 1/2] Fix double pendulum env --- ocean/double_pendulum/double_pendulum.h | 2 -- 1 file changed, 2 deletions(-) diff --git a/ocean/double_pendulum/double_pendulum.h b/ocean/double_pendulum/double_pendulum.h index dd2f411f6e..704d226a37 100644 --- a/ocean/double_pendulum/double_pendulum.h +++ b/ocean/double_pendulum/double_pendulum.h @@ -112,8 +112,6 @@ void c_reset(DoublePendulum* env) { env->episode_return = 0.0f; env->upright_steps = 0; env->max_upright_steps = 0; - env->rewards[0] = 0.0f; - env->terminals[0] = 0.0f; compute_observations(env); } From e4af14a520a9d58cc0576ea7f5d45b192c755435 Mon Sep 17 00:00:00 2001 From: Zilong Date: Wed, 19 Aug 2026 11:12:37 -0500 Subject: [PATCH 2/2] Add RK4 integration and a DeepMind-style reward to double pendulum Physics: replace the single semi-implicit Euler step with RK4 plus configurable substeps. Acceleration is factored out of integrate_physics into dp_accel so it can be evaluated at RK4 trial states; the dynamics themselves are unchanged. Measured energy drift over a 10s episode with zero force falls from 123% to 0.6355% at substeps=1 and 0.0006% at substeps=4. That matters here because a swing-up is entirely energy bookkeeping -- an integrator that leaks energy makes the catch easy for the wrong reason. Reward: replace 0.5*height + hold-streak bonus with a DeepMind Control Suite-style dense blend. The streak bookkeeping is kept because hold_time is derived from it, which keeps that metric reward-independent. On accurate physics the reward choice is decisive. Building on the terminal-signal fix in this branch's parent, at 500M steps from a dead hang, hold_time out of 600: Euler + streak 358 (previous behaviour) RK4 + streak 0 -- never catches RK4 + DeepMind 407 The streak reward pays 0.5100 for the first held step against 0.5000 for flying through the top, a 2% signal that Euler's energy drift was masking by acting as unintended dissipation. The DeepMind blend pays 4.2x at the same decision point. Note this removes the old reward rather than keeping it behind a flag, so existing double_pendulum results are not reproducible from here on. The standalone demo sets substeps and balance_bonus_weight explicitly. Its designated initializer would otherwise zero them, leaving the interactive binary on RK4x1 with w=0 while training runs RK4x4 with w=0.5. Also fix min_lr_ratio = 0 -> 0.0. The CLI infers each flag's type from its config default, so the int made --train.min-lr-ratio reject float values. Co-Authored-By: Claude Opus 5 --- config/double_pendulum.ini | 7 +- ocean/double_pendulum/binding.c | 3 + ocean/double_pendulum/double_pendulum.c | 2 + ocean/double_pendulum/double_pendulum.h | 108 +++++++++++++++++++----- 4 files changed, 98 insertions(+), 22 deletions(-) diff --git a/config/double_pendulum.ini b/config/double_pendulum.ini index fcb0cda022..f81d9e488e 100644 --- a/config/double_pendulum.ini +++ b/config/double_pendulum.ini @@ -15,6 +15,9 @@ link2_length = 0.5 gravity = 9.8 force_mag = 10.0 dt = 0.02 +substeps = 4 +# Blend between the DeepMind dense term and balance_quality. +balance_bonus_weight = 0.5 [policy] hidden_size = 128 @@ -29,10 +32,10 @@ decoder = DefaultDecoder [train] gpus = 1 seed = 42 -total_timesteps = 20000000 +total_timesteps = 500000000 learning_rate = 0.001 anneal_lr = 1 -min_lr_ratio = 0 +min_lr_ratio = 0.0 gamma = 0.99 gae_lambda = 0.95 replay_ratio = 1 diff --git a/ocean/double_pendulum/binding.c b/ocean/double_pendulum/binding.c index 5419f92445..60360fde8b 100644 --- a/ocean/double_pendulum/binding.c +++ b/ocean/double_pendulum/binding.c @@ -18,6 +18,9 @@ void my_init(Env* env, Dict* kwargs) { env->gravity = dict_get(kwargs, "gravity")->value; env->force_mag = dict_get(kwargs, "force_mag")->value; env->dt = dict_get(kwargs, "dt")->value; + env->substeps = (int)dict_get(kwargs, "substeps")->value; + env->balance_bonus_weight = dict_get(kwargs, "balance_bonus_weight")->value; + if (env->substeps < 1) env->substeps = 1; init(env); } diff --git a/ocean/double_pendulum/double_pendulum.c b/ocean/double_pendulum/double_pendulum.c index d3ab1ecf70..554364de96 100644 --- a/ocean/double_pendulum/double_pendulum.c +++ b/ocean/double_pendulum/double_pendulum.c @@ -21,6 +21,8 @@ int main(void) { .gravity = 9.8f, .force_mag = 10.0f, .dt = 0.02f, + .substeps = 4, + .balance_bonus_weight = 0.5f, }; init(&env); diff --git a/ocean/double_pendulum/double_pendulum.h b/ocean/double_pendulum/double_pendulum.h index 704d226a37..666962d47d 100644 --- a/ocean/double_pendulum/double_pendulum.h +++ b/ocean/double_pendulum/double_pendulum.h @@ -16,6 +16,15 @@ #define DP_HEIGHT 420 #define DP_SCALE 65.0f +// DeepMind Control Suite-style dense reward constants. +#define DP_CENTER_MARGIN 2.0f +#define DP_ANG_VEL_MARGIN 5.0f +#define DP_VELOCITY_FLOOR 0.60f +#define DP_LN_10 2.302585093f +#define DP_BALANCE_ANGLE_SCALE 0.30f +#define DP_BALANCE_ANG_VEL_SCALE 1.0f +#define DP_BALANCE_CART_SCALE 2.0f + typedef struct Log { float perf; float score; @@ -55,6 +64,8 @@ typedef struct DoublePendulum { float gravity; float force_mag; float dt; + int substeps; // physics substeps per control step (RK4) + float balance_bonus_weight; // blend between deepmind term and balance_quality } DoublePendulum; const Color PUFF_RED = (Color){187, 0, 0, 255}; @@ -152,16 +163,20 @@ static void solve_3x3(float A[3][3], float b[3], float x[3]) { x[2] = b[2]; } -void integrate_physics(DoublePendulum* env, float force) { +// Acceleration is a pure function of the angles, angular velocities and force +// (it does not depend on cart position/velocity), so RK4 can evaluate it at +// trial states. qdd = [xdd, th1dd, th2dd]. +static void dp_accel(DoublePendulum* env, const float th[2], const float w[2], + float force, float qdd[3]) { float m0 = env->cart_mass; float m1 = env->link1_mass; float m2 = env->link2_mass; float l1 = env->link1_length; float l2 = env->link2_length; - float t1 = env->theta1; - float t2 = env->theta2; - float w1 = env->theta1_dot; - float w2 = env->theta2_dot; + float t1 = th[0]; + float t2 = th[1]; + float w1 = w[0]; + float w2 = w[1]; float c1 = cosf(t1); float c2 = cosf(t2); float s1 = sinf(t1); @@ -179,22 +194,55 @@ void integrate_physics(DoublePendulum* env, float force) { (m1 + m2) * env->gravity * l1 * s1 - m2 * l1 * l2 * s12 * w2 * w2, m2 * env->gravity * l2 * s2 + m2 * l1 * l2 * s12 * w1 * w1, }; - float qdd[3]; solve_3x3(A, b, qdd); +} - env->x_dot += env->dt * qdd[0]; - env->theta1_dot += env->dt * qdd[1]; - env->theta2_dot += env->dt * qdd[2]; - env->x_dot = fminf(fmaxf(env->x_dot, -20.0f), 20.0f); - env->theta1_dot = fminf(fmaxf(env->theta1_dot, -30.0f), 30.0f); - env->theta2_dot = fminf(fmaxf(env->theta2_dot, -30.0f), 30.0f); - env->x += env->dt * env->x_dot; - env->theta1 = wrap_pi(env->theta1 + env->dt * env->theta1_dot); - env->theta2 = wrap_pi(env->theta2 + env->dt * env->theta2_dot); +// q = [x, theta1, theta2], v = [x_dot, theta1_dot, theta2_dot]. +static void dp_state_deriv(DoublePendulum* env, const float q[3], const float v[3], + float force, float dq[3], float dv[3]) { + for (int i = 0; i < 3; i++) dq[i] = v[i]; + float th[2] = {q[1], q[2]}; + float w[2] = {v[1], v[2]}; + dp_accel(env, th, w, force, dv); +} + +static void dp_rk4_step(DoublePendulum* env, float force, float h, + float q[3], float v[3]) { + float k1q[3], k1v[3], k2q[3], k2v[3], k3q[3], k3v[3], k4q[3], k4v[3]; + float tq[3], tv[3]; + dp_state_deriv(env, q, v, force, k1q, k1v); + for (int i = 0; i < 3; i++) { tq[i] = q[i] + 0.5f*h*k1q[i]; tv[i] = v[i] + 0.5f*h*k1v[i]; } + dp_state_deriv(env, tq, tv, force, k2q, k2v); + for (int i = 0; i < 3; i++) { tq[i] = q[i] + 0.5f*h*k2q[i]; tv[i] = v[i] + 0.5f*h*k2v[i]; } + dp_state_deriv(env, tq, tv, force, k3q, k3v); + for (int i = 0; i < 3; i++) { tq[i] = q[i] + h*k3q[i]; tv[i] = v[i] + h*k3v[i]; } + dp_state_deriv(env, tq, tv, force, k4q, k4v); + for (int i = 0; i < 3; i++) { + q[i] += (h/6.0f) * (k1q[i] + 2.0f*k2q[i] + 2.0f*k3q[i] + k4q[i]); + v[i] += (h/6.0f) * (k1v[i] + 2.0f*k2v[i] + 2.0f*k3v[i] + k4v[i]); + } +} + +void integrate_physics(DoublePendulum* env, float force) { + float q[3] = {env->x, env->theta1, env->theta2}; + float v[3] = {env->x_dot, env->theta1_dot, env->theta2_dot}; + int substeps = env->substeps > 0 ? env->substeps : 1; + float h = env->dt / (float)substeps; + for (int s = 0; s < substeps; s++) { + dp_rk4_step(env, force, h, q, v); + v[0] = fminf(fmaxf(v[0], -20.0f), 20.0f); + v[1] = fminf(fmaxf(v[1], -30.0f), 30.0f); + v[2] = fminf(fmaxf(v[2], -30.0f), 30.0f); + } + env->x = q[0]; + env->x_dot = v[0]; + env->theta1 = wrap_pi(q[1]); + env->theta1_dot = v[1]; + env->theta2 = wrap_pi(q[2]); + env->theta2_dot = v[2]; } float upright_reward(DoublePendulum* env, float force) { - (void)force; float tip_y = env->link1_length * cosf(env->theta1) + env->link2_length * cosf(env->theta2); float max_y = env->link1_length + env->link2_length; @@ -210,10 +258,30 @@ float upright_reward(DoublePendulum* env, float force) { env->max_upright_steps = env->upright_steps; } - float hold_bonus = fminf((float)env->upright_steps / 100.0f, 1.0f); - // Reward combines tip height and stable hold streak: 0.5*height + hold_bonus. - float reward = 0.5f * height + hold_bonus; - return fminf(fmaxf(reward, 0.0f), 1.5f); + // Reward: (1-w) * [h * centered * small_control * small_velocity] + w * balance_quality + float h = fminf(fmaxf(height, 0.0f), 1.0f); + float sw1 = env->theta1_dot / DP_ANG_VEL_MARGIN; + float sw2 = env->theta2_dot / DP_ANG_VEL_MARGIN; + float min_vel_tol = fminf(expf(-DP_LN_10 * sw1 * sw1), + expf(-DP_LN_10 * sw2 * sw2)); + float small_velocity = DP_VELOCITY_FLOOR + + (1.0f - DP_VELOCITY_FLOOR) * min_vel_tol; + float scaled_x = env->x / DP_CENTER_MARGIN; + float centered = 0.5f * (1.0f + expf(-DP_LN_10 * scaled_x * scaled_x)); + float a = env->force_mag > 0.0f ? force / env->force_mag : 0.0f; + float small_control = 0.2f * (4.0f + fmaxf(0.0f, 1.0f - a * a)); + float deepmind = h * centered * small_control * small_velocity; + float angle_mse = 0.5f * (env->theta1 * env->theta1 + + env->theta2 * env->theta2); + float ang_vel_mse = 0.5f * (env->theta1_dot * env->theta1_dot + + env->theta2_dot * env->theta2_dot); + float balance_x = env->x / DP_BALANCE_CART_SCALE; + float balance_quality = expf( + -angle_mse / (DP_BALANCE_ANGLE_SCALE * DP_BALANCE_ANGLE_SCALE) + -ang_vel_mse / (DP_BALANCE_ANG_VEL_SCALE * DP_BALANCE_ANG_VEL_SCALE) + -balance_x * balance_x); + float w = env->balance_bonus_weight; + return (1.0f - w) * deepmind + w * balance_quality; } void c_step(DoublePendulum* env) {