diff --git a/pufferlib/ocean/drive/drive.h b/pufferlib/ocean/drive/drive.h index 05f538dbab..af66c913bb 100644 --- a/pufferlib/ocean/drive/drive.h +++ b/pufferlib/ocean/drive/drive.h @@ -312,26 +312,27 @@ struct Drive { typedef struct { float min_val; float max_val; + int log_scale; } RewardBound; static const RewardBound REWARD_BOUNDS[NUM_REWARD_COEFS] = { - {2.0f, 12.0f}, // REWARD_COEF_GOAL_RADIUS δ_goal ~ U(2, 12) - {0.0f, 20.0f}, // REWARD_COEF_GOAL_SPEED - {0.0f, 3.0f}, // REWARD_COEF_COLLISION α_collision ~ U(0, 3) - {0.0f, 3.0f}, // REWARD_COEF_OFFROAD α_boundary ~ U(0, 3) - {0.0f, 0.1f}, // REWARD_COEF_COMFORT α_comfort ~ U(0, 0.1) - {2.5e-4f, 2.5e-2f}, // REWARD_COEF_LANE_ALIGN α_l-align ~ U(2.5e-4, 2.5e-2) - {0.0f, 1.0f}, // REWARD_COEF_VEL_ALIGN α_vel-align ~ U(0, 1) - {2.5e-4f, 7.5e-3f}, // REWARD_COEF_LANE_CENTER α_l-center ~ U(2.5e-4, 7.5e-3) - {-0.5f, 0.5f}, // REWARD_COEF_CENTER_BIAS α_center-bias ~ U(-0.5, 0.5) - {0.0f, 5e-3f}, // REWARD_COEF_VELOCITY α_velocity = 2.5e-3 (fixed) - {2.5e-4f, 7.5e-3f}, // REWARD_COEF_REVERSE α_reverse ~ U(2.5e-4, 7.5e-3) - {0.0f, 1.0f}, // REWARD_COEF_STOP_LINE α_stop-line ~ U(0, 1) - {0.0f, 5e-5f}, // REWARD_COEF_TIMESTEP α_timestep = 2.5e-5 (fixed) - {0.0f, 1.0f}, // REWARD_COEF_OVERSPEED - {0.8f, 1.25f}, // REWARD_COEF_THROTTLE C_throttle - {0.8f, 1.25f}, // REWARD_COEF_STEER C_steer - {0.666f, 1.5f}, // REWARD_COEF_ACC C_acc + {2.0f, 12.0f, 0}, // REWARD_COEF_GOAL_RADIUS δ_goal ~ U(2, 12) + {0.0f, 20.0f, 0}, // REWARD_COEF_GOAL_SPEED δ_goal-speed ~ U(0, 20) + {0.0f, 3.0f, 0}, // REWARD_COEF_COLLISION α_collision ~ U(0, 3) + {0.0f, 3.0f, 0}, // REWARD_COEF_OFFROAD α_boundary ~ U(0, 3) + {1e-5f, 0.1f, 1}, // REWARD_COEF_COMFORT α_comfort ~ logU(1e-5f, 0.1) + {2.5e-4f, 2.5e-2f, 1}, // REWARD_COEF_LANE_ALIGN α_l-align ~ logU(2.5e-4, 2.5e-2) + {0.0f, 1.0f, 0}, // REWARD_COEF_VEL_ALIGN α_vel-align ~ U(0, 1) + {2.5e-4f, 7.5e-3f, 1}, // REWARD_COEF_LANE_CENTER α_l-center ~ logU(2.5e-4, 7.5e-3) + {-0.5f, 0.5f, 0}, // REWARD_COEF_CENTER_BIAS α_center-bias ~ U(-0.5, 0.5) + {0.0f, 5e-3f, 0}, // REWARD_COEF_VELOCITY α_velocity = 2.5e-3 (fixed) + {2.5e-4f, 7.5e-3f, 1}, // REWARD_COEF_REVERSE α_reverse ~ logU(2.5e-4, 7.5e-3) + {0.0f, 1.0f, 0}, // REWARD_COEF_STOP_LINE α_stop-line ~ U(0, 1) + {0.0f, 5e-5f, 0}, // REWARD_COEF_TIMESTEP α_timestep = 2.5e-5 (fixed) + {0.0f, 1.0f, 0}, // REWARD_COEF_OVERSPEED α_overspeed ~ U(0, 1) + {0.8f, 1.25f, 0}, // REWARD_COEF_THROTTLE C_throttle + {0.8f, 1.25f, 0}, // REWARD_COEF_STEER C_steer + {0.666f, 1.5f, 0}, // REWARD_COEF_ACC C_acc }; // ======================================== @@ -381,16 +382,20 @@ static float compute_heading_diff(float heading1, float heading2) { return normalize_heading(heading1 - heading2); } -static float random_uniform(unsigned int *rng_state, float min_val, float max_val) { +static float sample_uniform(unsigned int *rng_state, float min_val, float max_val) { return min_val + ((float) rand_r(rng_state) / (float) RAND_MAX) * (max_val - min_val); } -static float mixed_uniform(unsigned int *rng_state, float a) { +static float sample_log_uniform(unsigned int *rng_state, float min_val, float max_val) { + return expf(sample_uniform(rng_state, logf(min_val), logf(max_val))); +} + +static float sample_mixed_uniform(unsigned int *rng_state, float a) { // Mixed uniform distribution X(a) = 0.5*U(1/a, 1) + 0.5*U(1, a) if ((float) rand_r(rng_state) / (float) RAND_MAX < 0.5f) { - return random_uniform(rng_state, 1.0f / a, 1.0f); + return sample_uniform(rng_state, 1.0f / a, 1.0f); } - return random_uniform(rng_state, 1.0f, a); + return sample_uniform(rng_state, 1.0f, a); } static void begin_episode_rng(Drive *env) { @@ -995,7 +1000,7 @@ static bool generate_new_goals_from_route(Drive *env, Agent *agent) { // Sample a spacing per goal, then walk the route placing goals at those forward distances. float goal_spacings_meters[MAX_GOALS]; for (int goal_idx = 0; goal_idx < env->num_goals; goal_idx++) { - goal_spacings_meters[goal_idx] = random_uniform(&env->rng_state, env->min_goal_spacing, env->max_goal_spacing); + goal_spacings_meters[goal_idx] = sample_uniform(&env->rng_state, env->min_goal_spacing, env->max_goal_spacing); } float goal_x[MAX_GOALS], goal_y[MAX_GOALS], goal_z[MAX_GOALS]; @@ -1113,7 +1118,7 @@ static bool generate_new_goals_from_map(Drive *env, Agent *agent) { int requested_goal_count = 1 + rand_r(&env->rng_state) % env->num_goals; float goal_spacings_meters[MAX_GOALS]; for (int goal_idx = 0; goal_idx < requested_goal_count; goal_idx++) { - goal_spacings_meters[goal_idx] = random_uniform(&env->rng_state, env->min_goal_spacing, env->max_goal_spacing); + goal_spacings_meters[goal_idx] = sample_uniform(&env->rng_state, env->min_goal_spacing, env->max_goal_spacing); } float goal_x[MAX_GOALS], goal_y[MAX_GOALS], goal_z[MAX_GOALS]; int goal_lane[MAX_GOALS]; @@ -1181,7 +1186,7 @@ static int roll_goals(Drive *env, Agent *agent) { } // Walk one spacing forward to find the appended goal before touching the window. - float spacing_meters = random_uniform(&env->rng_state, env->min_goal_spacing, env->max_goal_spacing); + float spacing_meters = sample_uniform(&env->rng_state, env->min_goal_spacing, env->max_goal_spacing); float next_x, next_y, next_z, next_s_on_lane; int next_lane_idx, next_cursor_idx; // next_cursor_idx unused: single-step append, no chaining if (!route_point_at_distance( @@ -2026,11 +2031,11 @@ static void add_log(Drive *env) { static inline void sample_erratic_flags(Drive *env, Agent *agent) { agent->is_blind_partner = (env->partner_blindness_prob > 0.0f - && random_uniform(&env->rng_state, 0.0f, 1.0f) < env->partner_blindness_prob) + && sample_uniform(&env->rng_state, 0.0f, 1.0f) < env->partner_blindness_prob) ? 1 : 0; agent->is_phantom_braker - = (env->phantom_braking_prob > 0.0f && random_uniform(&env->rng_state, 0.0f, 1.0f) < env->phantom_braking_prob) + = (env->phantom_braking_prob > 0.0f && sample_uniform(&env->rng_state, 0.0f, 1.0f) < env->phantom_braking_prob) ? 1 : 0; agent->phantom_braking_counter = 0; @@ -2054,14 +2059,15 @@ static void generate_reward_coefs(Drive *env, Agent *agent) { }; for (int i = 0; i < (int) (sizeof(random_coefs) / sizeof(random_coefs[0])); i++) { int c = random_coefs[i]; - agent->reward_coefs[c] - = random_uniform(&env->rng_state, REWARD_BOUNDS[c].min_val, REWARD_BOUNDS[c].max_val); + agent->reward_coefs[c] = REWARD_BOUNDS[c].log_scale + ? sample_log_uniform(&env->rng_state, REWARD_BOUNDS[c].min_val, REWARD_BOUNDS[c].max_val) + : sample_uniform(&env->rng_state, REWARD_BOUNDS[c].min_val, REWARD_BOUNDS[c].max_val); } agent->reward_coefs[REWARD_COEF_VELOCITY] = 2.5e-3f; agent->reward_coefs[REWARD_COEF_TIMESTEP] = 2.5e-5f; - agent->reward_coefs[REWARD_COEF_THROTTLE] = mixed_uniform(&env->rng_state, 1.25f); - agent->reward_coefs[REWARD_COEF_STEER] = mixed_uniform(&env->rng_state, 1.25f); - agent->reward_coefs[REWARD_COEF_ACC] = mixed_uniform(&env->rng_state, 1.5f); + agent->reward_coefs[REWARD_COEF_THROTTLE] = sample_mixed_uniform(&env->rng_state, 1.25f); + agent->reward_coefs[REWARD_COEF_STEER] = sample_mixed_uniform(&env->rng_state, 1.25f); + agent->reward_coefs[REWARD_COEF_ACC] = sample_mixed_uniform(&env->rng_state, 1.5f); } else { agent->reward_coefs[REWARD_COEF_GOAL_RADIUS] = env->goal_radius; agent->reward_coefs[REWARD_COEF_GOAL_SPEED] = env->goal_speed; @@ -2088,7 +2094,7 @@ static void generate_traffic_light_states(Drive *env) { float dt = env->dt; // 20% chance: disable ALL lights for this episode - int disable_all = (!env->eval_mode) && (random_uniform(&env->rng_state, 0.0f, 1.0f) < TL_EPISODE_DISABLE_PROB); + int disable_all = (!env->eval_mode) && (sample_uniform(&env->rng_state, 0.0f, 1.0f) < TL_EPISODE_DISABLE_PROB); for (int i = 0; i < env->num_traffic_elements; i++) { TrafficControlElement *tc = &env->traffic_elements[i]; @@ -2110,14 +2116,14 @@ static void generate_traffic_light_states(Drive *env) { if (!env->eval_mode) { // Individual removal - if (random_uniform(&env->rng_state, 0.0f, 1.0f) < TL_INDIVIDUAL_REMOVE_PROB) { + if (sample_uniform(&env->rng_state, 0.0f, 1.0f) < TL_INDIVIDUAL_REMOVE_PROB) { for (int t = 0; t < fill_steps; t++) { tc->states[t] = TRAFFIC_CONTROL_STATE_OFF; } continue; } // Always green - if (random_uniform(&env->rng_state, 0.0f, 1.0f) < TL_ALWAYS_GREEN_PROB) { + if (sample_uniform(&env->rng_state, 0.0f, 1.0f) < TL_ALWAYS_GREEN_PROB) { for (int t = 0; t < fill_steps; t++) { tc->states[t] = TRAFFIC_CONTROL_STATE_GREEN; } @@ -2132,12 +2138,12 @@ static void generate_traffic_light_states(Drive *env) { dur_yellow = TL_DEFAULT_YELLOW_DURATION; dur_red = TL_DEFAULT_RED_DURATION; } else { - dur_green = random_uniform(&env->rng_state, 0.1 * TL_DEFAULT_GREEN_DURATION, TL_DEFAULT_GREEN_DURATION); - dur_yellow = random_uniform( + dur_green = sample_uniform(&env->rng_state, 0.1 * TL_DEFAULT_GREEN_DURATION, TL_DEFAULT_GREEN_DURATION); + dur_yellow = sample_uniform( &env->rng_state, 0.5f * TL_DEFAULT_YELLOW_DURATION, 0.75f * TL_DEFAULT_YELLOW_DURATION); - dur_red = random_uniform(&env->rng_state, 0.15f * TL_DEFAULT_RED_DURATION, 5.0f * TL_DEFAULT_RED_DURATION); + dur_red = sample_uniform(&env->rng_state, 0.15f * TL_DEFAULT_RED_DURATION, 5.0f * TL_DEFAULT_RED_DURATION); } int steps_green = (int) (dur_green / dt); @@ -2258,12 +2264,12 @@ static bool spawn_agent(Drive *env, int agent_idx, int num_agents) { float spawn_length, spawn_width; if (env->eval_mode) { // Fixed size for eval mode - spawn_length = random_uniform(&env->rng_state, 2.0f, 5.5f); - spawn_width = random_uniform(&env->rng_state, 1.5f, 2.5f); + spawn_length = sample_uniform(&env->rng_state, 2.0f, 5.5f); + spawn_width = sample_uniform(&env->rng_state, 1.5f, 2.5f); } else { // Random size for training mode - spawn_length = random_uniform(&env->rng_state, 0.8f, 7.0f); - spawn_width = random_uniform(&env->rng_state, 0.8f, 2.7f); + spawn_length = sample_uniform(&env->rng_state, 0.8f, 7.0f); + spawn_width = sample_uniform(&env->rng_state, 0.8f, 2.7f); } if (spawn_width > spawn_length) { spawn_width = spawn_length; @@ -3592,8 +3598,17 @@ static int write_ego_obs(Drive *env, Agent *ego, float *obs, int obs_idx) { static int write_reward_target_obs(Drive *env, Agent *ego, float *obs, int obs_idx) { if (env->reward_conditioning) { for (int coef_idx = 0; coef_idx < NUM_REWARD_COEFS; coef_idx++) { - float normalized_coef = (ego->reward_coefs[coef_idx] - REWARD_BOUNDS[coef_idx].min_val) - / ((REWARD_BOUNDS[coef_idx].max_val - REWARD_BOUNDS[coef_idx].min_val) + 1e-8f); + float lo = REWARD_BOUNDS[coef_idx].min_val; + float hi = REWARD_BOUNDS[coef_idx].max_val; + float coef = ego->reward_coefs[coef_idx]; + float normalized_coef; + if (REWARD_BOUNDS[coef_idx].log_scale) { + // Match the log-uniform sampling so the conditioning signal stays even across [-1, 1]. + float clamped = fmaxf(lo, fminf(hi, coef)); + normalized_coef = (logf(clamped) - logf(lo)) / (logf(hi) - logf(lo)); + } else { + normalized_coef = (coef - lo) / ((hi - lo) + 1e-8f); + } float clamped_coef = fmaxf(0.0f, fminf(1.0f, normalized_coef)); obs[obs_idx++] = 2.0f * clamped_coef - 1.0f; } @@ -3622,7 +3637,7 @@ static int write_reward_target_obs(Drive *env, Agent *ego, float *obs, int obs_i } static int write_partner_obs(Drive *env, Agent *ego, int agent_idx, float *obs, int obs_idx, int *partner_count) { - if (ego->is_blind_partner && random_uniform(&env->rng_state, 0.0f, 1.0f) < env->partner_blindness_trigger_prob) { + if (ego->is_blind_partner && sample_uniform(&env->rng_state, 0.0f, 1.0f) < env->partner_blindness_trigger_prob) { int partner_obs_stride = env->obs_slots_partners_n * PARTNER_FEATURES; memset(&obs[obs_idx], 0, partner_obs_stride * sizeof(float)); *partner_count = 0; @@ -3990,7 +4005,7 @@ static void move_dynamics(Drive *env, int action_idx, int agent_idx) { phantom_braking_active = 1; } else if ( agent->is_phantom_braker && env->phantom_braking_trigger_prob > 0.0f - && random_uniform(&env->rng_state, 0.0f, 1.0f) < env->phantom_braking_trigger_prob) { + && sample_uniform(&env->rng_state, 0.0f, 1.0f) < env->phantom_braking_trigger_prob) { agent->phantom_braking_counter = env->phantom_braking_duration - 1; phantom_braking_active = 1; } @@ -4076,13 +4091,20 @@ static void move_dynamics(Drive *env, int action_idx, int agent_idx) { agent->jerk_lat = (new_a_lat - agent->accel_lat) / env->dt; agent->accel_long = new_a_long; agent->accel_lat = new_a_lat; - } else { - // DYNAMICS_MODEL_JERK dynamics model + } else if (env->dynamics_model == DYNAMICS_MODEL_JERK) { // Extract jerk action components float j_long, j_lat; - if (env->action_type == ACTION_TYPE_CONTINUOUS) { + if (env->action_type == ACTION_TYPE_DISCRETE) { + // Interpret action as a single integer: a = long_idx * num_lat + lat_idx + int *action_array = (int *) env->actions; + int num_lat = sizeof(JERK_LAT) / sizeof(JERK_LAT[0]); + int action_val = action_array[action_idx]; + int j_long_idx = action_val / num_lat; + int j_lat_idx = action_val % num_lat; + j_long = JERK_LONG[j_long_idx]; + j_lat = JERK_LAT[j_lat_idx]; + } else if (env->action_type == ACTION_TYPE_CONTINUOUS) { float (*action_array_f)[2] = (float (*)[2]) env->actions; - // Asymmetric scaling for longitudinal jerk to match discrete action space // Discrete: JERK_LONG = [-15, -4, 0, 4] (more braking than acceleration) float j_long_action = action_array_f[action_idx][0]; // [-1, 1] @@ -4091,18 +4113,8 @@ static void move_dynamics(Drive *env, int action_idx, int agent_idx) { } else { j_long = j_long_action * JERK_LONG[3]; // Positive: [0, 1] → [0, 4] (acceleration) } - // Symmetric scaling for lateral jerk j_lat = action_array_f[action_idx][1] * JERK_LAT[2]; - } else if (env->action_type == ACTION_TYPE_DISCRETE) { - // Interpret action as a single integer: a = long_idx * num_lat + lat_idx - int *action_array = (int *) env->actions; - int num_lat = sizeof(JERK_LAT) / sizeof(JERK_LAT[0]); - int action_val = action_array[action_idx]; - int j_long_idx = action_val / num_lat; - int j_lat_idx = action_val % num_lat; - j_long = JERK_LONG[j_long_idx]; - j_lat = JERK_LAT[j_lat_idx]; } if (phantom_braking_active) { diff --git a/tests/smoke_tests/data/drive_rollout_golden.json b/tests/smoke_tests/data/drive_rollout_golden.json index aabbbe509d..8e96952f46 100644 --- a/tests/smoke_tests/data/drive_rollout_golden.json +++ b/tests/smoke_tests/data/drive_rollout_golden.json @@ -6,7 +6,7 @@ "comfort_violation_count": 0.6436155717423622, "dnf_rate": 0.535904255319149, "episode_length": 13.48936170212766, - "episode_return": -1.14886574605678, + "episode_return": -0.7790729422518547, "lane_center_rate": 0.6959362321711601, "n": 16.0, "num_goals_reached": 0.00997340425531915, @@ -14,14 +14,14 @@ "red_light_violation_rate": 0.025930851063829786, "reward_components/ade": 0.0, "reward_components/collision": -0.042797685620632575, - "reward_components/comfort": -0.4419418668493311, + "reward_components/comfort": -0.09395423511716913, "reward_components/goal": 0.003324468085106383, - "reward_components/lane_align": -0.026303222145330398, - "reward_components/lane_center": -0.004736338347008333, + "reward_components/lane_align": -0.01126101781227725, + "reward_components/lane_center": -0.0026710387282172575, "reward_components/offroad": -0.6115261237037942, "reward_components/overspeed": 0.0, "reward_components/red_light": -0.014523565179688181, - "reward_components/reverse": -0.01046069894738971, + "reward_components/reverse": -0.005763038997835619, "reward_components/timestep": -8.58909574007873e-05, "reward_components/velocity": 0.00018516867541930356, "score": 0.0, diff --git a/tests/smoke_tests/data/drive_smoke_golden.json b/tests/smoke_tests/data/drive_smoke_golden.json index 694ef8bfc7..9dd4f2b4bf 100644 --- a/tests/smoke_tests/data/drive_smoke_golden.json +++ b/tests/smoke_tests/data/drive_smoke_golden.json @@ -1,48 +1,48 @@ { "env": { - "avg_distance_per_infraction": 11.854357493551154, - "avg_speed_per_agent": 1.3698952135286833, - "collision_rate": 0.01644736842105263, - "comfort_violation_count": 0.6486438763769049, - "dnf_rate": 0.5361842105263158, - "episode_length": 12.894736842105264, - "episode_return": -1.0695090011546486, - "lane_center_rate": 0.707680482613413, + "avg_distance_per_infraction": 12.15254513422648, + "avg_speed_per_agent": 1.369708173804813, + "collision_rate": 0.020833333333333332, + "comfort_violation_count": 0.6313128405147128, + "dnf_rate": 0.53125, + "episode_length": 13.38888888888889, + "episode_return": -0.747443911102083, + "lane_center_rate": 0.6955425375037723, "n": 16.0, - "num_goals_reached": 0.019736842105263157, + "num_goals_reached": 0.020833333333333332, "obs/max": 49.0, - "obs/mean": 0.24772183212917298, - "obs/min": -1.6471880599856377, - "offroad_rate": 0.40131578947368424, - "red_light_violation_rate": 0.03289473684210526, + "obs/mean": 0.24867265624925494, + "obs/min": -1.6757474672049284, + "offroad_rate": 0.4131944444444444, + "red_light_violation_rate": 0.024305555555555556, "reward_components/ade": 0.0, - "reward_components/collision": -0.027902447863628988, - "reward_components/comfort": -0.44484660343119975, + "reward_components/collision": -0.034726881318622164, + "reward_components/comfort": -0.09129529053138362, "reward_components/goal": 0.0, - "reward_components/lane_align": -0.026144957120873426, - "reward_components/lane_center": -0.004429814571171607, - "reward_components/offroad": -0.5440286923395959, + "reward_components/lane_align": -0.010500837311459085, + "reward_components/lane_center": -0.002869318226455814, + "reward_components/offroad": -0.593594147099389, "reward_components/overspeed": 0.0, - "reward_components/red_light": -0.012339777901376548, - "reward_components/reverse": -0.009904562618191304, - "reward_components/timestep": -8.252467025733112e-05, - "reward_components/velocity": 0.0001703243718733766, + "reward_components/red_light": -0.009579533297154639, + "reward_components/reverse": -0.0050971122044656016, + "reward_components/timestep": -8.596354139929947e-05, + "reward_components/velocity": 0.0003051760419718145, "score": 0.0, - "velocity_progress_sum": 0.017492458077245636 + "velocity_progress_sum": 0.028703066365172464 }, "losses": { - "approx_kl": 0.0013822934457233974, - "clipfrac": 0.0, - "ema_max": 2.47578889131546, - "entropy": 2.4796591826847623, - "explained_variance": -0.015610933303833008, - "filter_threshold": 0.024757888913154603, - "filtered_fraction": 0.081678956324447, - "kept_fraction": 0.918321043675553, - "masked_fraction": 0.13916015625, - "old_approx_kl": 0.0027940108307770322, - "policy_loss": -0.0008438232221773692, - "value_loss": 0.20083625614643097 + "approx_kl": 0.005358540724652509, + "clipfrac": 0.05141958842674891, + "ema_max": 2.412422329187393, + "entropy": 2.458675662676493, + "explained_variance": 0.009576380252838135, + "filter_threshold": 0.02412422329187393, + "filtered_fraction": 0.15625, + "kept_fraction": 0.84375, + "masked_fraction": 0.140625, + "old_approx_kl": 0.004483325717349847, + "policy_loss": -0.0015975118149071932, + "value_loss": 0.18829950441916785 }, "meta": { "bptt_horizon": 64,