Skip to content

Add log scaling support for reward coefficients #500

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
vcharraut merged 12 commits into from
Aug 4, 2026
+111 −99
Merged

Add log scaling support for reward coefficients #500

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
a09d608
Add log scaling support for reward coefficients and action types
vcharraut Jun 29, 2026
74fdb88
Update REWARD_COEF_COMFORT bounds for log scaling adjustment
vcharraut Jun 29, 2026
ea13493
Merge branch '3.0' into vcha/add-log-scaling
vcharraut Jun 29, 2026
6e84ae2
Merge branch '3.0' into vcha/add-log-scaling
vcharraut Jul 1, 2026
1deb3db
Merge 3.0
vcharraut Jul 16, 2026
5cab5cd
Merge 3.0
vcharraut Jul 30, 2026
0ac730a
Refactor dynamics model handling for discrete action type in move_dyn…
vcharraut Jul 30, 2026
361f725
Replace random_uniform with sample_uniform for goal spacing calculations
vcharraut Jul 30, 2026
2f3f094
Update golden JSON files with revised environment metrics and reward …
vcharraut Jul 30, 2026
a8762b2
Merge 3.0
vcharraut Aug 4, 2026
d975ae0
Refactor sampling functions to use consistent RNG state handling
vcharraut Aug 4, 2026
dcd1b31
Update golden JSON files with revised environment metrics and reward …
vcharraut Aug 4, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
130 changes: 71 additions & 59 deletions pufferlib/ocean/drive/drive.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
};

// ========================================
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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];
Expand Down Expand Up @@ -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];
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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;
Expand All @@ -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;
Expand All @@ -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];
Expand All @@ -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;
}
Expand All @@ -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);
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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;
}
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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;
}
Expand Down Expand Up @@ -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]
Expand All @@ -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) {
Expand Down
10 changes: 5 additions & 5 deletions tests/smoke_tests/data/drive_rollout_golden.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,22 +6,22 @@
"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,
"offroad_rate": 0.4115691489361702,
"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,
Expand Down
Loading
Loading