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 dd2f411f6e..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}; @@ -112,8 +123,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); } @@ -154,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); @@ -181,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; @@ -212,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) {