diff --git a/notebooks/notebook_utils.py b/notebooks/notebook_utils.py index 6db9e01d44..a58c17ce64 100644 --- a/notebooks/notebook_utils.py +++ b/notebooks/notebook_utils.py @@ -69,8 +69,8 @@ "reward_randomization": False, "target_type": "static", "map_dir": MAP_DIR, - "collision_behavior": 1, - "offroad_behavior": 1, + "collision_behavior": "stop", + "offroad_behavior": "stop", "obs_slots_lane_n": 80, "obs_slots_boundary_n": 80, "obs_lane_stride": 1, diff --git a/pufferlib/config/ocean/drive.ini b/pufferlib/config/ocean/drive.ini index e3acd64d65..bac4f3dd34 100644 --- a/pufferlib/config/ocean/drive.ini +++ b/pufferlib/config/ocean/drive.ini @@ -58,12 +58,12 @@ dynamics_model = "jerk" dt = 0.1 ; Optional nonzero launch speed for gigaflow random spawns spawn_initial_speed = 0.0 -; Collision behavior - options: 0 - Ignore, 1 - Stop, 2 - Remove -collision_behavior = 1 -; Offroad behavior - options: 0 - Ignore, 1 - Stop, 2 - Remove -offroad_behavior = 1 -; Traffic light behavior - options: 0 - Ignore, 1 - Stop, 2 - Remove -traffic_light_behavior = 1 +; Collision behavior - options: "ignore", "stop", "remove" +collision_behavior = "stop" +; Offroad behavior - options: "ignore", "stop", "remove" +offroad_behavior = "stop" +; Traffic light behavior - options: "ignore", "stop", "remove" +traffic_light_behavior = "stop" ; Share static map geometry (roads/grid/lane-graph) across envs using the same map - 0 off, 1 on use_map_cache = 0 ; Number of steps before reset @@ -269,10 +269,10 @@ interval = 250 mode = "inline" clean = true env.eval_mode = 1 -env.collision_behavior = 1 -env.offroad_behavior = 1 -; Explicit 0 wins over the clean macro's red-light enforcement. -env.traffic_light_behavior = 0 +env.collision_behavior = "stop" +env.offroad_behavior = "stop" +; Explicit "ignore" wins over the clean macro's red-light enforcement. +env.traffic_light_behavior = "ignore" env.reward_randomization = False env.termination_mode = 0 env.num_agents = 1024 diff --git a/pufferlib/config/puffer_drive.yaml b/pufferlib/config/puffer_drive.yaml index 641d6f1826..3883169ab8 100644 --- a/pufferlib/config/puffer_drive.yaml +++ b/pufferlib/config/puffer_drive.yaml @@ -54,12 +54,12 @@ env: dt: 0.1 # Optional nonzero launch speed for gigaflow random spawns spawn_initial_speed: 0.0 - # Collision behavior - options: 0 - Ignore, 1 - Stop, 2 - Remove - collision_behavior: 1 - # Offroad behavior - options: 0 - Ignore, 1 - Stop, 2 - Remove - offroad_behavior: 1 - # Traffic light behavior - options: 0 - Ignore, 1 - Stop, 2 - Remove - traffic_light_behavior: 1 + # Collision behavior - options: "ignore", "stop", "remove" + collision_behavior: stop + # Offroad behavior - options: "ignore", "stop", "remove" + offroad_behavior: stop + # Traffic light behavior - options: "ignore", "stop", "remove" + traffic_light_behavior: stop # Share static map geometry (roads/grid/lane-graph) across envs using the same map - 0 off, 1 on use_map_cache: 0 # Number of steps before reset @@ -306,10 +306,10 @@ eval: clean: 'true' env: eval_mode: 1 - collision_behavior: 1 - offroad_behavior: 1 - # Explicit 0 wins over the clean macro's red-light enforcement. - traffic_light_behavior: 0 + collision_behavior: stop + offroad_behavior: stop + # Explicit "ignore" wins over the clean macro's red-light enforcement. + traffic_light_behavior: ignore reward_randomization: false termination_mode: 0 num_agents: 1024 diff --git a/pufferlib/config_schema.py b/pufferlib/config_schema.py new file mode 100644 index 0000000000..49aefb8155 --- /dev/null +++ b/pufferlib/config_schema.py @@ -0,0 +1,173 @@ +"""Structured-config schemas for load_config — Hydra migration phase 2. + +Each entry in ENV_SCHEMAS is an OmegaConf structured schema merged over the +composed monolithic YAML in `pufferlib.pufferl.load_config`, so type errors, +unknown keys, and invalid enum names fail at config load instead of deep in +env construction. + +Every field is MISSING on purpose: the YAML is the single source of values, +the schema only contributes types. A key present here but absent from the +YAML surfaces as a missing-value error at load time. + +Enum int values mirror the #defines in pufferlib/ocean/drive/drive.h (the +source of truth, exported to Python via binding). Only the *names* gate +config validation — drive.py still maps names to binding constants itself. + +Naming convention (Google enum style): every C #define is +`_`, e.g. +`InfractionBehavior.stop` <-> `INFRACTION_BEHAVIOR_STOP`. This makes the C +name mechanically derivable from the Python one (and vice versa) instead of +needing a lookup table. NonVehicleController is the one exception: it +reuses Controller's C constants (`CONTROLLER_*`) since it's the same +underlying controller, plus a Python-only `auto` sentinel that drive.py +resolves before it ever reaches C. tests/unit_tests/test_config_schema.py +walks every class here and asserts the derived name exists in binding with +the matching value, so a new enum is checked automatically without a new +hand-written assert. +""" + +from dataclasses import dataclass +from enum import Enum + +from omegaconf import MISSING + + +class SimulationMode(Enum): + gigaflow = 0 + replay = 1 + + +class ActionType(Enum): + discrete = 0 + continuous = 1 + + +class DynamicsModel(Enum): + classic = 0 + jerk = 1 + + +class InfractionBehavior(Enum): + ignore = 0 + stop = 1 + remove = 2 + + +class ControlMode(Enum): + control_vehicles = 0 + control_agents = 1 + control_wosac = 2 + control_sdc_only = 3 + + +class Controller(Enum): + static = 0 + policy = 1 + replay = 2 + idm = 3 + + +class NonVehicleController(Enum): + # "auto" is config-side only: drive.py resolves it to a Controller + # (replay when non_sdc_controller is idm, else non_sdc_controller). + auto = -1 + static = 0 + policy = 1 + replay = 2 + idm = 3 + + +class InitMode(Enum): + create_all_valid = 0 + create_only_controlled = 1 + + +class TargetType(Enum): + static = 0 + dynamic = 1 + + +@dataclass +class DriveEnvConfig: + simulation_mode: SimulationMode = MISSING + num_agents: int = MISSING + min_agents_per_env: int = MISSING + max_agents_per_env: int = MISSING + action_type: ActionType = MISSING + dynamics_model: DynamicsModel = MISSING + dt: float = MISSING + spawn_initial_speed: float = MISSING + collision_behavior: InfractionBehavior = MISSING + offroad_behavior: InfractionBehavior = MISSING + traffic_light_behavior: InfractionBehavior = MISSING + use_map_cache: int = MISSING + scenario_length: int = MISSING + resample_frequency: int = MISSING + termination_mode: int = MISSING + inactive_agent_threshold: float = MISSING + terminate_on_goal: int = MISSING + init_step: int = MISSING + init_step_spread: bool = MISSING + init_step_min_horizon: int = MISSING + control_mode: ControlMode = MISSING + sdc_controller: Controller = MISSING + non_sdc_controller: Controller = MISSING + non_vehicle_controller: NonVehicleController = MISSING + init_mode: InitMode = MISSING + compute_eval_metrics: bool = MISSING + target_type: TargetType = MISSING + goal_on_lane: bool = MISSING + goal_radius: float = MISSING + goal_speed: float = MISSING + num_goals: int = MISSING + min_goal_spacing: float = MISSING + max_goal_spacing: float = MISSING + reward_conditioning: bool = MISSING + reward_randomization: bool = MISSING + reward_goal: float = MISSING + reward_collision: float = MISSING + reward_offroad: float = MISSING + reward_stop_line: float = MISSING + reward_comfort: float = MISSING + reward_lane_align: float = MISSING + reward_vel_align: float = MISSING + reward_lane_center: float = MISSING + reward_center_bias: float = MISSING + reward_velocity: float = MISSING + reward_reverse: float = MISSING + reward_timestep: float = MISSING + reward_overspeed: float = MISSING + reward_ade: float = MISSING + map_dir: str = MISSING + num_maps: int = MISSING + obs_slots_lane_n: int = MISSING + obs_slots_boundary_n: int = MISSING + obs_slots_partners_n: int = MISSING + obs_slots_traffic_controls_n: int = MISSING + obs_dropout_lane: float = MISSING + obs_dropout_boundary: float = MISSING + obs_lane_stride: int = MISSING + obs_boundary_stride: int = MISSING + obs_norm_goal_offset_m: float = MISSING + obs_norm_xy_offset_m: float = MISSING + obs_norm_veh_length_m: float = MISSING + obs_norm_veh_width_m: float = MISSING + obs_norm_road_seg_length_m: float = MISSING + obs_norm_road_seg_width_m: float = MISSING + obs_range_road_front_m: float = MISSING + obs_range_road_behind_m: float = MISSING + obs_range_road_side_m: float = MISSING + obs_range_partner_m: float = MISSING + obs_range_traffic_control_m: float = MISSING + partner_blindness_prob: float = MISSING + partner_blindness_trigger_prob: float = MISSING + phantom_braking_prob: float = MISSING + phantom_braking_trigger_prob: float = MISSING + phantom_braking_duration: int = MISSING + + +# env_name -> structured schema for the `env` section. Envs without an entry +# load unvalidated (phase-1 behavior). +ENV_SCHEMAS = { + "puffer_drive": DriveEnvConfig, +} diff --git a/pufferlib/ocean/benchmark/manager.py b/pufferlib/ocean/benchmark/manager.py index 93eb49093d..9308841ed0 100644 --- a/pufferlib/ocean/benchmark/manager.py +++ b/pufferlib/ocean/benchmark/manager.py @@ -43,7 +43,7 @@ "partner_blindness_prob": 0.0, "phantom_braking_prob": 0.0, "phantom_braking_trigger_prob": 0.0, - "traffic_light_behavior": 1, + "traffic_light_behavior": "stop", } diff --git a/pufferlib/ocean/drive/binding.c b/pufferlib/ocean/drive/binding.c index d7a2356150..7d91df39f4 100644 --- a/pufferlib/ocean/drive/binding.c +++ b/pufferlib/ocean/drive/binding.c @@ -1745,7 +1745,7 @@ static PyObject *my_shared(PyObject *self, PyObject *args, PyObject *kwargs) { PyErr_SetString(PyExc_ValueError, "min_agents_per_env must be <= max_agents_per_env"); return NULL; } - if (simulation_mode == SIMULATION_GIGAFLOW && num_agents < min_agents_per_env) { + if (simulation_mode == SIMULATION_MODE_GIGAFLOW && num_agents < min_agents_per_env) { PyErr_SetString(PyExc_ValueError, "num_agents must be >= min_agents_per_env"); return NULL; } @@ -1753,7 +1753,7 @@ static PyObject *my_shared(PyObject *self, PyObject *args, PyObject *kwargs) { srand(seed); // GIGAFLOW mode: use random sampling for agent counts per env - if (simulation_mode == SIMULATION_GIGAFLOW) { + if (simulation_mode == SIMULATION_MODE_GIGAFLOW) { if (eval_mode) { // Eval mode: fixed agent count, sequential map cycling int agents_per_env = max_agents_per_env; diff --git a/pufferlib/ocean/drive/drive.c b/pufferlib/ocean/drive/drive.c index 1c9365ca22..cf12d8d530 100644 --- a/pufferlib/ocean/drive/drive.c +++ b/pufferlib/ocean/drive/drive.c @@ -22,7 +22,7 @@ void test_drivenet() { // Weights* weights = load_weights("resources/drive/puffer_drive_weights.bin"); Weights *weights = load_weights("puffer_drive_weights.bin"); - DriveNet *net = init_drivenet(weights, num_agents, CLASSIC); + DriveNet *net = init_drivenet(weights, num_agents, DYNAMICS_MODEL_CLASSIC); forward(net, observations, actions); for (int i = 0; i < num_agents * num_actions; i++) { diff --git a/pufferlib/ocean/drive/drive.h b/pufferlib/ocean/drive/drive.h index 61fb7a746b..99892a972a 100644 --- a/pufferlib/ocean/drive/drive.h +++ b/pufferlib/ocean/drive/drive.h @@ -55,14 +55,14 @@ #define EGO_IDX 0 // Initialization modes -#define INIT_ALL_VALID 0 -#define INIT_ONLY_CONTROLLABLE_AGENTS 1 +#define INIT_MODE_CREATE_ALL_VALID 0 +#define INIT_MODE_CREATE_ONLY_CONTROLLED 1 // Control modes -#define CONTROL_VEHICLES 0 -#define CONTROL_AGENTS 1 -#define CONTROL_WOSAC 2 -#define CONTROL_SDC_ONLY 3 +#define CONTROL_MODE_VEHICLES 0 +#define CONTROL_MODE_AGENTS 1 +#define CONTROL_MODE_WOSAC 2 +#define CONTROL_MODE_SDC_ONLY 3 // Controller modes #define CONTROLLER_STATIC 0 @@ -71,8 +71,12 @@ #define CONTROLLER_IDM 3 // Simulation modes -#define SIMULATION_GIGAFLOW 0 -#define SIMULATION_REPLAY 1 +#define SIMULATION_MODE_GIGAFLOW 0 +#define SIMULATION_MODE_REPLAY 1 + +// Action types +#define ACTION_TYPE_DISCRETE 0 +#define ACTION_TYPE_CONTINUOUS 1 // Lane selection scoring #define LANE_SELECTION_DISTANCE_WEIGHT 0.7f @@ -109,8 +113,9 @@ #define MULTI_LANE_HALF_SCORE_TIME 5.7f // seconds // Collision/Infraction behaviors -#define STOP_AGENT 1 -#define REMOVE_AGENT 2 +#define INFRACTION_BEHAVIOR_IGNORE 0 +#define INFRACTION_BEHAVIOR_STOP 1 +#define INFRACTION_BEHAVIOR_REMOVE 2 #define MAX_SPEED 40.0f @@ -124,8 +129,8 @@ #define ROAD_QUERY_ENTITY_COUNT (MAX_ENTITIES_PER_CELL * 25) // TARGET_TYPE modes (controls what target info is in observations) -#define TARGET_STATIC 0 -#define TARGET_DYNAMIC 1 +#define TARGET_TYPE_STATIC 0 +#define TARGET_TYPE_DYNAMIC 1 // Observation feature counts #define EGO_FEATURES 10 @@ -162,19 +167,19 @@ static const int ROAD_OFFSETS[25][2] {1, 1}, {2, 1}, {-2, 2}, {-1, 2}, {0, 2}, {1, 2}, {2, 2}}; // Dynamics Models -#define CLASSIC 0 -#define JERK 1 +#define DYNAMICS_MODEL_CLASSIC 0 +#define DYNAMICS_MODEL_JERK 1 static const float ACCEL_LONG_LIMIT[2] = {-5.0f, 2.5f}; static const float ACCEL_LAT_LIMIT[2] = {-4.0f, 4.0f}; #define STEERING_LIMIT 0.667f static const float REAR_AXLE_RATIO = 0.5f; -// Jerk action space (for JERK dynamics model) +// Jerk action space (for DYNAMICS_MODEL_JERK dynamics model) static const float JERK_LONG[4] = {-15.0f, -4.0f, 0.0f, 4.0f}; static const float JERK_LAT[3] = {-4.0f, 0.0f, 4.0f}; -// Classic action space (for CLASSIC dynamics model) +// Classic action space (for DYNAMICS_MODEL_CLASSIC dynamics model) static const float ACCELERATION_VALUES[7] = {-4.0000f, -2.6670f, -1.3330f, -0.0000f, 1.3330f, 2.6670f, 4.0000f}; static const float STEERING_VALUES[9] = {-0.667f, -0.500f, -0.333f, -0.167f, 0.000f, 0.167f, 0.333f, 0.500f, 0.667f}; @@ -397,10 +402,10 @@ struct Drive { int target_type; int goal_on_lane; char *ini_file; - int collision_behavior; // 0 = none, 1=stop, 2 = remove - int offroad_behavior; // 0 = none, 1=stop, 2 = remove - int traffic_light_behavior; // 0 = none, 1=stop, 2 = remove - int use_map_cache; // 0 = each env owns its map copy, 1 = share static geometry across envs + int collision_behavior; // INFRACTION_BEHAVIOR_IGNORE, INFRACTION_BEHAVIOR_STOP, or INFRACTION_BEHAVIOR_REMOVE + int offroad_behavior; // INFRACTION_BEHAVIOR_IGNORE, INFRACTION_BEHAVIOR_STOP, or INFRACTION_BEHAVIOR_REMOVE + int traffic_light_behavior; // INFRACTION_BEHAVIOR_IGNORE, INFRACTION_BEHAVIOR_STOP, or INFRACTION_BEHAVIOR_REMOVE + int use_map_cache; // 0 = each env owns its map copy, 1 = share static geometry across envs struct SharedMapData *shared_map; // non-NULL when this env borrows cached geometry // Metadata fields char scenario_id[128]; @@ -615,9 +620,9 @@ static inline void update_agent_radius(Agent *agent) { } static inline void apply_infraction_behavior(Agent *agent, int behavior) { - if (behavior == STOP_AGENT && !agent->stopped) { + if (behavior == INFRACTION_BEHAVIOR_STOP && !agent->stopped) { agent->stopped = 1; - } else if (behavior == REMOVE_AGENT && !agent->removed) { + } else if (behavior == INFRACTION_BEHAVIOR_REMOVE && !agent->removed) { agent->removed = 1; } } @@ -1875,7 +1880,7 @@ static int compute_new_route(Drive *env, int agent_idx, int current_lane_idx) { int num_goals = env->num_goals; float min_route_distance; // NOTE: make both multipliers config values and tune from a metric (route regenerations per 1k env steps). - if (env->target_type == TARGET_STATIC) { + if (env->target_type == TARGET_TYPE_STATIC) { min_route_distance = env->max_goal_spacing * num_goals * 2.0f; } else { min_route_distance = env->min_goal_spacing * num_goals * 20.0f; @@ -2086,7 +2091,7 @@ static void compute_goals(Drive *env, int agent_idx) { // If we reached the end of the current path, compute a new route and retry. // Bounded by iter <= 4 to prevent infinite loops on degenerate maps. if (needed_s >= path_end_s) { - if (env->simulation_mode == SIMULATION_GIGAFLOW) { + if (env->simulation_mode == SIMULATION_MODE_GIGAFLOW) { if (iter > 3) { printf("[GIGAFLOW WARNING] -> Max iterations in compute_goals for agent %d\n", agent_idx); agent->removed = 1; @@ -3389,7 +3394,7 @@ static void set_start_position(Drive *env) { Agent *agent = &env->agents[i]; // Initialize simulation trajectory from logged trajectory at init_step - if (env->simulation_mode == SIMULATION_REPLAY) { + if (env->simulation_mode == SIMULATION_MODE_REPLAY) { // Clamp init_step to ensure we don't go out of bounds int step = env->init_step; if (step >= agent->trajectory_size) { @@ -3456,11 +3461,11 @@ static bool should_control_agent(Drive *env, int agent_idx) { Agent *agent = &env->agents[agent_idx]; - if (env->control_mode == CONTROL_SDC_ONLY) { + if (env->control_mode == CONTROL_MODE_SDC_ONLY) { return agent_idx == EGO_IDX && agent->route_length != 0; } - if (env->control_mode == CONTROL_WOSAC) { + if (env->control_mode == CONTROL_MODE_WOSAC) { for (int j = 0; j < env->num_tracks_to_predict; j++) { if (env->tracks_to_predict[j] == agent_idx) { return true; @@ -3471,9 +3476,9 @@ static bool should_control_agent(Drive *env, int agent_idx) { // Standard mode: check type, distance to goal, and expert status bool type_is_controllable = false; - if (env->control_mode == CONTROL_VEHICLES) { + if (env->control_mode == CONTROL_MODE_VEHICLES) { type_is_controllable = (agent->type == VEHICLE); - } else { // CONTROL_AGENTS mode + } else { // CONTROL_MODE_AGENTS mode type_is_controllable = is_controllable_agent(agent->type); } @@ -3482,7 +3487,7 @@ static bool should_control_agent(Drive *env, int agent_idx) { } // In REPLAY mode without route data, control agents spawning far enough from their goal - if (env->simulation_mode == SIMULATION_REPLAY && agent->route_length == 0) { + if (env->simulation_mode == SIMULATION_MODE_REPLAY && agent->route_length == 0) { float dx = agent->current_goal_x - agent->log_trajectory_x[env->init_step]; float dy = agent->current_goal_y - agent->log_trajectory_y[env->init_step]; float dz = agent->current_goal_z - agent->log_trajectory_z[env->init_step]; @@ -3523,7 +3528,7 @@ void set_active_agents(Drive *env) { env->num_agents = 0; // Total agents created // In GIGAFLOW mode, spawn agents dynamically on the map - if (env->simulation_mode == SIMULATION_GIGAFLOW) { + if (env->simulation_mode == SIMULATION_MODE_GIGAFLOW) { int num_agents_to_create = env->num_controllable_agents; // Initialize agents for GIGAFLOW mode @@ -3561,7 +3566,7 @@ void set_active_agents(Drive *env) { } // In REPLAY mode, determine which agents to control - bool is_log_replay = (env->control_mode == CONTROL_SDC_ONLY); + bool is_log_replay = (env->control_mode == CONTROL_MODE_SDC_ONLY); // In log-replay mode, no cap on actors int max_agents = is_log_replay ? env->num_total_agents : env->num_max_agents; @@ -3582,9 +3587,9 @@ void set_active_agents(Drive *env) { bool should_create = false; if (is_log_replay) { should_create = true; // Log-replay: all valid agents - } else if (env->init_mode == INIT_ALL_VALID) { + } else if (env->init_mode == INIT_MODE_CREATE_ALL_VALID) { should_create = true; // All valid entities - } else if (env->control_mode == CONTROL_VEHICLES) { + } else if (env->control_mode == CONTROL_MODE_VEHICLES) { should_create = (agent->type == VEHICLE); } else { // Control all agents should_create = (is_controllable_agent(agent->type)); @@ -3604,7 +3609,7 @@ void set_active_agents(Drive *env) { env->active_agent_count++; env->agents[i].active_agent = 1; env->agents[i].controller = resolve_agent_controller(env, i, 1, 0); - } else if (is_log_replay || env->init_mode != INIT_ONLY_CONTROLLABLE_AGENTS) { + } else if (is_log_replay || env->init_mode != INIT_MODE_CREATE_ONLY_CONTROLLED) { static_agent_indices[env->static_agent_count] = i; env->static_agent_count++; env->agents[i].active_agent = 0; @@ -3648,11 +3653,11 @@ void set_active_agents(Drive *env) { } void move_expert(Drive *env, int agent_idx) { - if (env->simulation_mode == SIMULATION_GIGAFLOW) { + if (env->simulation_mode == SIMULATION_MODE_GIGAFLOW) { printf("[GIGAFLOW ERROR] -> move_expert() called in GIGAFLOW mode\n"); return; } - bool is_log_replay = (env->control_mode == CONTROL_SDC_ONLY); + bool is_log_replay = (env->control_mode == CONTROL_MODE_SDC_ONLY); Agent *agent = &env->agents[agent_idx]; int t = env->timestep; @@ -3691,7 +3696,7 @@ void move_expert(Drive *env, int agent_idx) { } void remove_bad_trajectories(Drive *env) { - if (env->control_mode == CONTROL_WOSAC) { + if (env->control_mode == CONTROL_MODE_WOSAC) { return; // Leave all trajectories in WOSAC control mode } @@ -3860,7 +3865,7 @@ void init(Drive *env) { env->road_dropout_enabled = (env->obs_slots_lane_kept < env->obs_slots_lane_n) || (env->obs_slots_boundary_kept < env->obs_slots_boundary_n); env->logs_capacity = 0; - if (env->simulation_mode == SIMULATION_GIGAFLOW) { + if (env->simulation_mode == SIMULATION_MODE_GIGAFLOW) { int steps = env->scenario_length; if (steps > 0) { for (int i = 0; i < env->num_traffic_elements; i++) { @@ -3886,13 +3891,13 @@ void init(Drive *env) { } set_active_agents(env); env->logs_capacity = env->active_agent_count; - if (env->simulation_mode == SIMULATION_REPLAY) { + if (env->simulation_mode == SIMULATION_MODE_REPLAY) { remove_bad_trajectories(env); } set_start_position(env); env->logs = (Log *) calloc(env->active_agent_count, sizeof(Log)); - if (env->simulation_mode == SIMULATION_REPLAY) { + if (env->simulation_mode == SIMULATION_MODE_REPLAY) { for (int i = 0; i < env->active_agent_count; i++) { int agent_idx = env->active_agent_indices[i]; Agent *agent = &env->agents[agent_idx]; @@ -3978,9 +3983,9 @@ void c_close(Drive *env) { } static int compute_observation_size(Drive *env) { - int target_features = (env->target_type == TARGET_STATIC) ? STATIC_TARGET_FEATURES - : (env->target_type == TARGET_DYNAMIC) ? DYNAMIC_TARGET_FEATURES - : 0; + int target_features = (env->target_type == TARGET_TYPE_STATIC) ? STATIC_TARGET_FEATURES + : (env->target_type == TARGET_TYPE_DYNAMIC) ? DYNAMIC_TARGET_FEATURES + : 0; return EGO_FEATURES + PARTNER_FEATURES * env->obs_slots_partners_n + ROAD_FEATURES * (env->obs_slots_lane_kept + env->obs_slots_boundary_kept) + TRAFFIC_CONTROL_FEATURES * env->obs_slots_traffic_controls_n + OBS_VALID_COUNT_FEATURES @@ -4148,7 +4153,7 @@ static void compute_metrics(Drive *env, int agent_idx, int log_idx) { } // Compute log-replay metrics - if (env->simulation_mode == SIMULATION_REPLAY) { + if (env->simulation_mode == SIMULATION_MODE_REPLAY) { // Compute displacement error float displacement_error = compute_displacement_error(agent, env->timestep); if (displacement_error > 0.0f) { // Only count valid displacements @@ -4471,7 +4476,7 @@ static void compute_rewards(Drive *env, int i) { // Goal reward if (agent->metrics_array[REACHED_GOAL_IDX] > 0.0f) { float weight = 1.0f; - if (env->simulation_mode == SIMULATION_GIGAFLOW) { + if (env->simulation_mode == SIMULATION_MODE_GIGAFLOW) { if (agent->current_goal_idx == env->num_goals && agent->sim_speed > agent->reward_coefs[REWARD_COEF_GOAL_SPEED]) { weight = 0.0f; @@ -4627,7 +4632,7 @@ static int write_reward_target_obs(Drive *env, Agent *ego, float *obs, int obs_i } } - if (env->target_type == TARGET_STATIC) { + if (env->target_type == TARGET_TYPE_STATIC) { for (int wp_idx = 0; wp_idx < env->num_goals; wp_idx++) { if (wp_idx < ego->current_goal_idx) { obs[obs_idx++] = 0.0f; @@ -4649,7 +4654,7 @@ static int write_reward_target_obs(Drive *env, Agent *ego, float *obs, int obs_i return obs_idx; } - if (env->target_type == TARGET_DYNAMIC && ego->path != NULL && ego->path->num_waypoints > 0) { + if (env->target_type == TARGET_TYPE_DYNAMIC && ego->path != NULL && ego->path->num_waypoints > 0) { for (int wp_idx = 0; wp_idx < env->num_goals; wp_idx++) { int clamped_wp_idx = fmin(ego->closest_path_idx_wp + wp_idx, ego->path->num_waypoints - 1); if (clamped_wp_idx < 0) { @@ -5011,12 +5016,12 @@ static void move_dynamics(Drive *env, int action_idx, int agent_idx) { phantom_braking_active = 1; } - if (env->dynamics_model == CLASSIC) { + if (env->dynamics_model == DYNAMICS_MODEL_CLASSIC) { // Classic dynamics model float acceleration = 0.0f; float steering = 0.0f; - if (env->action_type == 0) { // discrete + if (env->action_type == ACTION_TYPE_DISCRETE) { // Interpret action as a single integer: a = accel_idx * num_steer + steer_idx int *action_array = (int *) env->actions; int num_steer = sizeof(STEERING_VALUES) / sizeof(STEERING_VALUES[0]); @@ -5025,7 +5030,7 @@ static void move_dynamics(Drive *env, int action_idx, int agent_idx) { int steering_index = action_val % num_steer; acceleration = ACCELERATION_VALUES[acceleration_index]; steering = STEERING_VALUES[steering_index]; - } else if (env->action_type == 1) { // continuous + } else if (env->action_type == ACTION_TYPE_CONTINUOUS) { float (*action_array_f)[2] = (float (*)[2]) env->actions; acceleration = action_array_f[action_idx][0]; steering = action_array_f[action_idx][1]; @@ -5093,10 +5098,10 @@ static void move_dynamics(Drive *env, int action_idx, int agent_idx) { agent->accel_long = new_a_long; agent->accel_lat = new_a_lat; } else { - // JERK dynamics model + // DYNAMICS_MODEL_JERK dynamics model // Extract jerk action components float j_long, j_lat; - if (env->action_type == 1) { // continuous + 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 @@ -5110,7 +5115,7 @@ static void move_dynamics(Drive *env, int action_idx, int agent_idx) { // Symmetric scaling for lateral jerk j_lat = action_array_f[action_idx][1] * JERK_LAT[2]; - } else if (env->action_type == 0) { // discrete + } 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]); @@ -5255,7 +5260,7 @@ void c_reset(Drive *env) { env->timestep = env->init_step; - if (env->simulation_mode == SIMULATION_GIGAFLOW) { + if (env->simulation_mode == SIMULATION_MODE_GIGAFLOW) { generate_traffic_light_states(env); int num_reset = 0; for (int x = 0; x < env->active_agent_count; x++) { @@ -5310,7 +5315,7 @@ void c_reset(Drive *env) { sample_erratic_flags(env, agent); generate_reward_coefs(env, agent); - if (env->simulation_mode == SIMULATION_REPLAY) { + if (env->simulation_mode == SIMULATION_MODE_REPLAY) { int start = env->init_step > 0 ? env->init_step : 0; int remaining = agent->trajectory_size - 1 - start; if (remaining < 1) { @@ -5368,7 +5373,7 @@ void c_step(Drive *env) { Agent *agent = &env->agents[background_idx]; if (agent->controller == CONTROLLER_IDM) { move_idm(env, background_idx); - } else if (agent->controller == CONTROLLER_REPLAY && env->simulation_mode == SIMULATION_REPLAY) { + } else if (agent->controller == CONTROLLER_REPLAY && env->simulation_mode == SIMULATION_MODE_REPLAY) { move_expert(env, background_idx); } } @@ -5382,7 +5387,7 @@ void c_step(Drive *env) { move_dynamics(env, i, agent_idx); } else if (agent->controller == CONTROLLER_IDM) { move_idm(env, agent_idx); - } else if (agent->controller == CONTROLLER_REPLAY && env->simulation_mode == SIMULATION_REPLAY) { + } else if (agent->controller == CONTROLLER_REPLAY && env->simulation_mode == SIMULATION_MODE_REPLAY) { move_expert(env, agent_idx); } } @@ -5435,8 +5440,8 @@ void c_step(Drive *env) { } } - if (env->terminate_on_goal == 1 && env->simulation_mode == SIMULATION_REPLAY - && env->control_mode == CONTROL_SDC_ONLY) { + if (env->terminate_on_goal == 1 && env->simulation_mode == SIMULATION_MODE_REPLAY + && env->control_mode == CONTROL_MODE_SDC_ONLY) { for (int i = 0; i < env->active_agent_count; i++) { Agent *agent = &env->agents[env->active_agent_indices[i]]; if (agent->metrics_array[REACHED_GOAL_IDX] > 0.0f && agent->current_goal_idx == env->num_goals) { @@ -5465,7 +5470,7 @@ void c_step(Drive *env) { if (agent->metrics_array[REACHED_GOAL_IDX] > 0.0f && agent->current_goal_idx == env->num_goals) { // Last goal reached env->logs[i].num_goals_reached += 1; - if (env->simulation_mode == SIMULATION_REPLAY) { + if (env->simulation_mode == SIMULATION_MODE_REPLAY) { // Replay mode: leave current_goal_idx saturated so the // reached-goal condition won't fire again. Re-generating // route-based goals on WOMD maps fails (removed=1). diff --git a/pufferlib/ocean/drive/drive.py b/pufferlib/ocean/drive/drive.py index 2914dc2ace..52f85fa104 100644 --- a/pufferlib/ocean/drive/drive.py +++ b/pufferlib/ocean/drive/drive.py @@ -45,9 +45,9 @@ def __init__( max_goal_spacing=60.0, num_goals=3, goal_radius=2.0, - collision_behavior=0, - offroad_behavior=0, - traffic_light_behavior=0, + collision_behavior="ignore", + offroad_behavior="ignore", + traffic_light_behavior="ignore", use_map_cache=0, # emit_completed_episodes=True: env emits one summary dict per # completed episode via info (drained from a per-env C-side queue). @@ -149,15 +149,27 @@ def __init__( self.num_goals = num_goals self.target_type_str = target_type if target_type == "static": - self.target_type = binding.TARGET_STATIC + self.target_type = binding.TARGET_TYPE_STATIC elif target_type == "dynamic": - self.target_type = binding.TARGET_DYNAMIC + self.target_type = binding.TARGET_TYPE_DYNAMIC else: raise ValueError(f"target_type must be 'static' or 'dynamic'. Got: {target_type}") self.goal_on_lane = int(bool(goal_on_lane)) - self.collision_behavior = collision_behavior - self.offroad_behavior = offroad_behavior - self.traffic_light_behavior = traffic_light_behavior + infraction_behavior_values = { + "ignore": binding.INFRACTION_BEHAVIOR_IGNORE, + "stop": binding.INFRACTION_BEHAVIOR_STOP, + "remove": binding.INFRACTION_BEHAVIOR_REMOVE, + } + for behavior_name, behavior in ( + ("collision_behavior", collision_behavior), + ("offroad_behavior", offroad_behavior), + ("traffic_light_behavior", traffic_light_behavior), + ): + if behavior not in infraction_behavior_values: + raise ValueError(f"{behavior_name} must be one of 'ignore', 'stop', or 'remove'. Got: {behavior}") + self.collision_behavior = infraction_behavior_values[collision_behavior] + self.offroad_behavior = infraction_behavior_values[offroad_behavior] + self.traffic_light_behavior = infraction_behavior_values[traffic_light_behavior] if use_map_cache not in (0, 1): raise ValueError(f"use_map_cache must be 0 (off) or 1 (on). Got: {use_map_cache}") self.use_map_cache = use_map_cache @@ -171,9 +183,9 @@ def __init__( self.resample_frequency = resample_frequency self.dynamics_model = dynamics_model if dynamics_model == "classic": - self.dynamics_model_flag = 0 + self.dynamics_model_flag = binding.DYNAMICS_MODEL_CLASSIC elif dynamics_model == "jerk": - self.dynamics_model_flag = 1 + self.dynamics_model_flag = binding.DYNAMICS_MODEL_JERK else: raise ValueError(f"dynamics_model must be 'classic' or 'jerk'. Got: {dynamics_model}") self.eval_mode = eval_mode @@ -274,14 +286,14 @@ def __init__( self.map_files = sorted(os.path.join(map_dir, f) for f in os.listdir(map_dir) if f.endswith(".bin")) if self.simulation_mode_str == "gigaflow": - self.simulation_mode = 0 + self.simulation_mode = binding.SIMULATION_MODE_GIGAFLOW elif self.simulation_mode_str == "replay": - self.simulation_mode = 1 + self.simulation_mode = binding.SIMULATION_MODE_REPLAY else: raise ValueError(f"simulation_mode must be one of 'gigaflow' or 'replay'. Got: {self.simulation_mode_str}") if self.init_step_spread: - if self.simulation_mode != 1: + if self.simulation_mode != binding.SIMULATION_MODE_REPLAY: raise ValueError( "init_step_spread is only supported in replay simulation_mode (it seeds each environment at a different expert timestep)." ) @@ -291,13 +303,13 @@ def __init__( ) if self.control_mode_str == "control_vehicles": - self.control_mode = 0 + self.control_mode = binding.CONTROL_MODE_VEHICLES elif self.control_mode_str == "control_agents": - self.control_mode = 1 + self.control_mode = binding.CONTROL_MODE_AGENTS elif self.control_mode_str == "control_wosac": - self.control_mode = 2 + self.control_mode = binding.CONTROL_MODE_WOSAC elif self.control_mode_str == "control_sdc_only": - self.control_mode = 3 + self.control_mode = binding.CONTROL_MODE_SDC_ONLY else: raise ValueError( "control_mode must be one of 'control_vehicles', 'control_agents', 'control_wosac', or " @@ -332,16 +344,16 @@ def __init__( self.non_vehicle_controller = controller_values[self.non_vehicle_controller_str] if self.init_mode_str == "create_all_valid": - self.init_mode = 0 + self.init_mode = binding.INIT_MODE_CREATE_ALL_VALID elif self.init_mode_str == "create_only_controlled": - self.init_mode = 1 + self.init_mode = binding.INIT_MODE_CREATE_ONLY_CONTROLLED else: raise ValueError( f"init_mode must be one of 'create_all_valid' or 'create_only_controlled'. Got: {self.init_mode_str}" ) if action_type == "discrete": - self._action_type_flag = 0 + self._action_type_flag = binding.ACTION_TYPE_DISCRETE if dynamics_model == "classic": # Joint action space (assume dependence) self.single_action_space = gymnasium.spaces.MultiDiscrete([7 * 9]) @@ -353,7 +365,7 @@ def __init__( else: raise ValueError(f"dynamics_model must be 'classic' or 'jerk'. Got: {dynamics_model}") elif action_type == "continuous": - self._action_type_flag = 1 + self._action_type_flag = binding.ACTION_TYPE_CONTINUOUS self.single_action_space = gymnasium.spaces.Box(low=-1, high=1, shape=(2,), dtype=np.float32) else: raise ValueError(f"action_space must be 'discrete' or 'continuous'. Got: {action_type}") diff --git a/pufferlib/ocean/drive/drivenet.h b/pufferlib/ocean/drive/drivenet.h index 939c0afa39..375e506537 100644 --- a/pufferlib/ocean/drive/drivenet.h +++ b/pufferlib/ocean/drive/drivenet.h @@ -53,11 +53,11 @@ DriveNet *init_drivenet(Weights *weights, int num_agents, int dynamics_model) { // Determine action space size based on dynamics model int action_size, logit_sizes[2]; int action_dim; - if (dynamics_model == CLASSIC) { + if (dynamics_model == DYNAMICS_MODEL_CLASSIC) { action_size = 7 * 9; // Joint action space (7 accel × 9 steering = 63) logit_sizes[0] = 7 * 9; action_dim = 1; - } else { // JERK + } else { // DYNAMICS_MODEL_JERK action_size = 4 * 3; // Joint action space (4 longitudinal × 3 lateral = 12) logit_sizes[0] = 4 * 3; action_dim = 1; diff --git a/pufferlib/ocean/drive/render.h b/pufferlib/ocean/drive/render.h index 6211ce0247..9f8affd9e0 100644 --- a/pufferlib/ocean/drive/render.h +++ b/pufferlib/ocean/drive/render.h @@ -735,8 +735,8 @@ void draw_agent_obs(Drive *env, int agent_index, int mode, int obs_only, int las int ego_dim = EGO_FEATURES; int num_reward_coefs = env->reward_conditioning ? NUM_REWARD_COEFS : 0; - int target_features = (env->target_type == TARGET_STATIC) ? env->num_goals * STATIC_TARGET_FEATURES - : env->num_goals * DYNAMIC_TARGET_FEATURES; + int target_features = (env->target_type == TARGET_TYPE_STATIC) ? env->num_goals * STATIC_TARGET_FEATURES + : env->num_goals * DYNAMIC_TARGET_FEATURES; int max_obs = compute_observation_size(env); float (*observations)[max_obs] = (float (*)[max_obs]) env->observations; float *agent_obs = &observations[agent_index][0]; @@ -1432,7 +1432,7 @@ void draw_scene(Drive *env, Client *client, int mode, int obs_only, int lasers, EndMode3D(); // Draw track indices for the tracks to predict - if (mode == 1 && env->control_mode == CONTROL_WOSAC) { + if (mode == 1 && env->control_mode == CONTROL_MODE_WOSAC) { float map_height = env->grid_map->top_left_y - env->grid_map->bottom_right_y; float pixels_per_world_unit = client->height / map_height; diff --git a/pufferlib/ocean/env_binding.h b/pufferlib/ocean/env_binding.h index 579965df94..ca06e510aa 100644 --- a/pufferlib/ocean/env_binding.h +++ b/pufferlib/ocean/env_binding.h @@ -1402,12 +1402,27 @@ PyMODINIT_FUNC PyInit_binding(void) { PyModule_AddIntConstant(m, "SCORE_F32_FIELDS", SCORE_F32_FIELDS); PyModule_AddIntConstant(m, "TRAFFIC_I16_FIELDS", TRAFFIC_I16_FIELDS); PyModule_AddIntConstant(m, "NUM_REWARD_COEFS", NUM_REWARD_COEFS); - PyModule_AddIntConstant(m, "TARGET_STATIC", TARGET_STATIC); - PyModule_AddIntConstant(m, "TARGET_DYNAMIC", TARGET_DYNAMIC); + PyModule_AddIntConstant(m, "TARGET_TYPE_STATIC", TARGET_TYPE_STATIC); + PyModule_AddIntConstant(m, "TARGET_TYPE_DYNAMIC", TARGET_TYPE_DYNAMIC); PyModule_AddIntConstant(m, "CONTROLLER_STATIC", CONTROLLER_STATIC); PyModule_AddIntConstant(m, "CONTROLLER_POLICY", CONTROLLER_POLICY); PyModule_AddIntConstant(m, "CONTROLLER_REPLAY", CONTROLLER_REPLAY); PyModule_AddIntConstant(m, "CONTROLLER_IDM", CONTROLLER_IDM); + PyModule_AddIntConstant(m, "INFRACTION_BEHAVIOR_IGNORE", INFRACTION_BEHAVIOR_IGNORE); + PyModule_AddIntConstant(m, "INFRACTION_BEHAVIOR_STOP", INFRACTION_BEHAVIOR_STOP); + PyModule_AddIntConstant(m, "INFRACTION_BEHAVIOR_REMOVE", INFRACTION_BEHAVIOR_REMOVE); + PyModule_AddIntConstant(m, "SIMULATION_MODE_GIGAFLOW", SIMULATION_MODE_GIGAFLOW); + PyModule_AddIntConstant(m, "SIMULATION_MODE_REPLAY", SIMULATION_MODE_REPLAY); + PyModule_AddIntConstant(m, "ACTION_TYPE_DISCRETE", ACTION_TYPE_DISCRETE); + PyModule_AddIntConstant(m, "ACTION_TYPE_CONTINUOUS", ACTION_TYPE_CONTINUOUS); + PyModule_AddIntConstant(m, "DYNAMICS_MODEL_CLASSIC", DYNAMICS_MODEL_CLASSIC); + PyModule_AddIntConstant(m, "DYNAMICS_MODEL_JERK", DYNAMICS_MODEL_JERK); + PyModule_AddIntConstant(m, "CONTROL_MODE_VEHICLES", CONTROL_MODE_VEHICLES); + PyModule_AddIntConstant(m, "CONTROL_MODE_AGENTS", CONTROL_MODE_AGENTS); + PyModule_AddIntConstant(m, "CONTROL_MODE_WOSAC", CONTROL_MODE_WOSAC); + PyModule_AddIntConstant(m, "CONTROL_MODE_SDC_ONLY", CONTROL_MODE_SDC_ONLY); + PyModule_AddIntConstant(m, "INIT_MODE_CREATE_ALL_VALID", INIT_MODE_CREATE_ALL_VALID); + PyModule_AddIntConstant(m, "INIT_MODE_CREATE_ONLY_CONTROLLED", INIT_MODE_CREATE_ONLY_CONTROLLED); PyObject_SetAttrString(m, "MULTI_LANE_FULL_SCORE_TIME", PyFloat_FromDouble(MULTI_LANE_FULL_SCORE_TIME)); PyObject_SetAttrString(m, "MULTI_LANE_HALF_SCORE_TIME", PyFloat_FromDouble(MULTI_LANE_HALF_SCORE_TIME)); diff --git a/pufferlib/ocean/env_config.h b/pufferlib/ocean/env_config.h index 9a0f454ec1..19cd273175 100644 --- a/pufferlib/ocean/env_config.h +++ b/pufferlib/ocean/env_config.h @@ -74,6 +74,23 @@ typedef struct { int phantom_braking_duration; } env_init_config; +// Shared "ignore"/"stop"/"remove" enum for the collision/offroad/traffic-light +// behavior keys. Values mirror INFRACTION_BEHAVIOR_IGNORE/INFRACTION_BEHAVIOR_STOP/INFRACTION_BEHAVIOR_REMOVE in +// drive.h. +static int parse_infraction_behavior(const char *name, const char *value) { + if (strcmp(value, "\"ignore\"") == 0 || strcmp(value, "ignore") == 0) { + return 0; // INFRACTION_BEHAVIOR_IGNORE + } + if (strcmp(value, "\"stop\"") == 0 || strcmp(value, "stop") == 0) { + return 1; // INFRACTION_BEHAVIOR_STOP + } + if (strcmp(value, "\"remove\"") == 0 || strcmp(value, "remove") == 0) { + return 2; // INFRACTION_BEHAVIOR_REMOVE + } + fprintf(stderr, "Invalid %s value '%s': must be \"ignore\", \"stop\", or \"remove\"\n", name, value); + exit(1); +} + // INI file parser handler - parses all environment configuration from drive.ini static int handler(void *config, const char *section, const char *name, const char *value) { env_init_config *env_config = (env_init_config *) config; @@ -81,35 +98,35 @@ static int handler(void *config, const char *section, const char *name, const ch if (MATCH("env", "action_type")) { if (strcmp(value, "\"discrete\"") == 0 || strcmp(value, "discrete") == 0) { - env_config->action_type = 0; // DISCRETE + env_config->action_type = 0; // ACTION_TYPE_DISCRETE } else if (strcmp(value, "\"continuous\"") == 0 || strcmp(value, "continuous") == 0) { - env_config->action_type = 1; // CONTINUOUS + env_config->action_type = 1; // ACTION_TYPE_CONTINUOUS } else { - printf("Warning: Unknown action_type value '%s', defaulting to DISCRETE\n", value); - env_config->action_type = 0; // Default to DISCRETE + printf("Warning: Unknown action_type value '%s', defaulting to ACTION_TYPE_DISCRETE\n", value); + env_config->action_type = 0; // Default to ACTION_TYPE_DISCRETE } } else if (MATCH("env", "dynamics_model")) { if (strcmp(value, "\"classic\"") == 0 || strcmp(value, "classic") == 0) { - env_config->dynamics_model = 0; // CLASSIC + env_config->dynamics_model = 0; // DYNAMICS_MODEL_CLASSIC } else if (strcmp(value, "\"jerk\"") == 0 || strcmp(value, "jerk") == 0) { - env_config->dynamics_model = 1; // JERK + env_config->dynamics_model = 1; // DYNAMICS_MODEL_JERK } else { - printf("Warning: Unknown dynamics_model value '%s', defaulting to JERK\n", value); - env_config->dynamics_model = 1; // Default to JERK + printf("Warning: Unknown dynamics_model value '%s', defaulting to DYNAMICS_MODEL_JERK\n", value); + env_config->dynamics_model = 1; // Default to DYNAMICS_MODEL_JERK } } else if (MATCH("env", "collision_behavior")) { - env_config->collision_behavior = atoi(value); + env_config->collision_behavior = parse_infraction_behavior(name, value); } else if (MATCH("env", "offroad_behavior")) { - env_config->offroad_behavior = atoi(value); + env_config->offroad_behavior = parse_infraction_behavior(name, value); } else if (MATCH("env", "traffic_light_behavior")) { - env_config->traffic_light_behavior = atoi(value); + env_config->traffic_light_behavior = parse_infraction_behavior(name, value); } else if (MATCH("env", "use_map_cache")) { env_config->use_map_cache = atoi(value); } else if (MATCH("env", "target_type")) { if (strcmp(value, "\"static\"") == 0 || strcmp(value, "static") == 0) { - env_config->target_type = 0; // TARGET_STATIC + env_config->target_type = 0; // TARGET_TYPE_STATIC } else if (strcmp(value, "\"dynamic\"") == 0 || strcmp(value, "dynamic") == 0) { - env_config->target_type = 1; // TARGET_DYNAMIC + env_config->target_type = 1; // TARGET_TYPE_DYNAMIC } else { printf("Warning: Unknown target_type value '%s', defaulting to static\n", value); env_config->target_type = 0; diff --git a/pufferlib/pufferl.py b/pufferlib/pufferl.py index 9e1ebb26d6..6663298225 100644 --- a/pufferlib/pufferl.py +++ b/pufferlib/pufferl.py @@ -41,6 +41,7 @@ import pufferlib.vector import pufferlib.pytorch import pufferlib.viz +from pufferlib.config_schema import ENV_SCHEMAS try: @@ -2118,10 +2119,19 @@ def load_config(env_name, config_dir=None): with initialize_config_dir(config_dir=config_dir, version_base=None): cfg = compose(config_name=env_name, overrides=overrides) + # Structured-schema validation (types, enum names, unknown keys) for envs + # that declare one. Overrides are already composed in, so CLI typos fail + # here too — at load time, not deep in env construction. + env_schema = ENV_SCHEMAS.get(env_name) + if env_schema is not None: + cfg["env"] = OmegaConf.merge(OmegaConf.structured(env_schema), cfg["env"]) + # Plain nested dict — the contract every downstream consumer relies on. # Protein's sweep.suggest() writes arbitrary keys into it, so no - # struct-mode OmegaConf objects may leak past this point. - args = defaultdict(dict, OmegaConf.to_container(cfg, resolve=True)) + # struct-mode OmegaConf objects may leak past this point. enum_to_str + # converts validated enum members back to their names; throw_on_missing + # rejects schema keys the YAML no longer provides. + args = defaultdict(dict, OmegaConf.to_container(cfg, resolve=True, enum_to_str=True, throw_on_missing=True)) args["train"]["use_rnn"] = args["rnn_name"] is not None diff --git a/pufferlib/viz.py b/pufferlib/viz.py index fbc89a46d6..ab563a2811 100644 --- a/pufferlib/viz.py +++ b/pufferlib/viz.py @@ -537,7 +537,7 @@ def unpack_obs( obs_flat = obs_flat[None, :] if isinstance(target_type, int): - target_type = "static" if target_type == binding.TARGET_STATIC else "dynamic" + target_type = "static" if target_type == binding.TARGET_TYPE_STATIC else "dynamic" ego_dim = binding.EGO_FEATURES @@ -632,7 +632,7 @@ def plot_observation( target_type: 0 for goal only, 1 for waypoints only, 2 for both """ if isinstance(target_type, int): - target_type = "static" if target_type == binding.TARGET_STATIC else "dynamic" + target_type = "static" if target_type == binding.TARGET_TYPE_STATIC else "dynamic" fig, ax = plt.subplots(figsize=(20, 20)) @@ -707,7 +707,7 @@ def plot_observation( s = 100 ax.scatter(wp_x, wp_y, color=color, marker=marker, s=s, zorder=15) - # Add dynamics info text for JERK model + # Add dynamics info text for DYNAMICS_MODEL_JERK model ego_info = f"Speed: {ego_speed:.2f}\nLane Centering: {lcenter:.2f}\nLane Align: {lalign:.2f}\nSpeed Limit: {speed_limit:.2f}" ego_info += f"\nSteering: {steering_angle:.3f}\naccel_long: {accel_long:.2f}\naccel_lat: {accel_lat:.2f}" diff --git a/scripts/render_scenario.py b/scripts/render_scenario.py index b929a0a04b..e79be6fc7a 100644 --- a/scripts/render_scenario.py +++ b/scripts/render_scenario.py @@ -167,8 +167,8 @@ def main(): env_overrides["max_agents_per_env"] = num_agents elif cli.simulation_mode == "replay": # Don't stop the SDC for offroad/collision so the full trajectory renders - env_overrides["offroad_behavior"] = 0 - env_overrides["collision_behavior"] = 0 + env_overrides["offroad_behavior"] = "ignore" + env_overrides["collision_behavior"] = "ignore" env_overrides["scenario_length"] = steps env_overrides["resample_frequency"] = steps diff --git a/tests/drive/include/drive_fixture.h b/tests/drive/include/drive_fixture.h index 6d34e75e37..e1912b84c8 100644 --- a/tests/drive/include/drive_fixture.h +++ b/tests/drive/include/drive_fixture.h @@ -53,7 +53,7 @@ static inline Drive drive_test_env_config( env.render_mode = RENDER_WINDOW; snprintf(env.resource_root, sizeof(env.resource_root), "%s", DRIVE_TEST_REPO_ROOT "/pufferlib/resources/drive"); env.action_type = 0; - env.dynamics_model = CLASSIC; + env.dynamics_model = DYNAMICS_MODEL_CLASSIC; env.reward_goal = 1.0f; env.reward_collision = 3.0f; env.reward_offroad = 3.0f; @@ -68,9 +68,9 @@ static inline Drive drive_test_env_config( env.reward_timestep = 0.000025f; env.reward_overspeed = 0.05f; env.reward_ade = 0.0f; - env.collision_behavior = 0; - env.offroad_behavior = 0; - env.traffic_light_behavior = 0; + env.collision_behavior = INFRACTION_BEHAVIOR_IGNORE; + env.offroad_behavior = INFRACTION_BEHAVIOR_IGNORE; + env.traffic_light_behavior = INFRACTION_BEHAVIOR_IGNORE; env.use_map_cache = use_map_cache; env.emit_completed_episodes = 1; env.goal_radius = 2.0f; @@ -78,7 +78,7 @@ static inline Drive drive_test_env_config( env.min_goal_spacing = 20.0f; env.max_goal_spacing = 60.0f; env.num_goals = 3; - env.target_type = TARGET_STATIC; + env.target_type = TARGET_TYPE_STATIC; env.goal_on_lane = 1; env.obs_slots_lane_n = 32; env.obs_slots_boundary_n = 32; @@ -97,8 +97,8 @@ static inline Drive drive_test_env_config( env.num_max_agents = 64; env.init_step = 0; env.timestep = 0; - env.init_mode = INIT_ALL_VALID; - env.control_mode = simulation_mode == SIMULATION_REPLAY ? CONTROL_SDC_ONLY : CONTROL_VEHICLES; + env.init_mode = INIT_MODE_CREATE_ALL_VALID; + env.control_mode = simulation_mode == SIMULATION_MODE_REPLAY ? CONTROL_MODE_SDC_ONLY : CONTROL_MODE_VEHICLES; env.sdc_controller = CONTROLLER_POLICY; env.non_sdc_controller = CONTROLLER_POLICY; env.non_vehicle_controller = CONTROLLER_REPLAY; @@ -137,7 +137,7 @@ static inline void drive_set_neutral_actions(Drive *env) { if (env->action_type == 0) { int *actions = (int *) env->actions; int neutral; - if (env->dynamics_model == JERK) { + if (env->dynamics_model == DYNAMICS_MODEL_JERK) { int num_long = sizeof(JERK_LONG) / sizeof(JERK_LONG[0]); int num_lat = sizeof(JERK_LAT) / sizeof(JERK_LAT[0]); neutral = (num_long / 2) * num_lat + (num_lat / 2); diff --git a/tests/drive/test_drive_dynamics.c b/tests/drive/test_drive_dynamics.c index fc6ea3bab9..85c8a49cd2 100644 --- a/tests/drive/test_drive_dynamics.c +++ b/tests/drive/test_drive_dynamics.c @@ -3,9 +3,9 @@ static int test_classic_action_clipping(void) { srand(19); - Drive env = drive_test_make_env(drive_carla_map(), SIMULATION_GIGAFLOW, 1, 0); + Drive env = drive_test_make_env(drive_carla_map(), SIMULATION_MODE_GIGAFLOW, 1, 0); env.action_type = 1; - env.dynamics_model = CLASSIC; + env.dynamics_model = DYNAMICS_MODEL_CLASSIC; Agent *agent = &env.agents[env.active_agent_indices[0]]; agent->sim_speed_signed = MAX_SPEED * 0.9f; agent->steering_angle = STEERING_LIMIT * 0.9f; @@ -20,9 +20,9 @@ static int test_classic_action_clipping(void) { static int test_jerk_action_clipping(void) { srand(23); - Drive env = drive_test_make_env(drive_carla_map(), SIMULATION_GIGAFLOW, 1, 0); + Drive env = drive_test_make_env(drive_carla_map(), SIMULATION_MODE_GIGAFLOW, 1, 0); env.action_type = 1; - env.dynamics_model = JERK; + env.dynamics_model = DYNAMICS_MODEL_JERK; Agent *agent = &env.agents[env.active_agent_indices[0]]; agent->sim_speed_signed = MAX_SPEED * 0.9f; agent->steering_angle = STEERING_LIMIT * 0.9f; @@ -70,9 +70,9 @@ static int test_dynamics_removed_agent_invalidated(void) { static int test_neutral_actions_zero_out(void) { { - Drive env = drive_test_make_env(drive_carla_map(), SIMULATION_GIGAFLOW, 1, 0); + Drive env = drive_test_make_env(drive_carla_map(), SIMULATION_MODE_GIGAFLOW, 1, 0); env.action_type = 0; - env.dynamics_model = CLASSIC; + env.dynamics_model = DYNAMICS_MODEL_CLASSIC; drive_set_neutral_actions(&env); int action_val = ((int *) env.actions)[0]; int num_steer = sizeof(STEERING_VALUES) / sizeof(STEERING_VALUES[0]); @@ -83,9 +83,9 @@ static int test_neutral_actions_zero_out(void) { free_allocated(&env); } { - Drive env = drive_test_make_env(drive_carla_map(), SIMULATION_GIGAFLOW, 1, 0); + Drive env = drive_test_make_env(drive_carla_map(), SIMULATION_MODE_GIGAFLOW, 1, 0); env.action_type = 0; - env.dynamics_model = JERK; + env.dynamics_model = DYNAMICS_MODEL_JERK; drive_set_neutral_actions(&env); int action_val = ((int *) env.actions)[0]; int num_lat = sizeof(JERK_LAT) / sizeof(JERK_LAT[0]); diff --git a/tests/drive/test_drive_env_smoke.c b/tests/drive/test_drive_env_smoke.c index 3477e659f1..d02c77d43d 100644 --- a/tests/drive/test_drive_env_smoke.c +++ b/tests/drive/test_drive_env_smoke.c @@ -40,20 +40,20 @@ static int run_case(const char *name, const char *map_file, int simulation_mode, } static int test_carla_gigaflow_load_step_log(void) { - return run_case("carla-gigaflow", drive_carla_map(), SIMULATION_GIGAFLOW, 32); + return run_case("carla-gigaflow", drive_carla_map(), SIMULATION_MODE_GIGAFLOW, 32); } static int test_nuplan_gigaflow_load_step_log(void) { - return run_case("nuplan-gigaflow", drive_nuplan_map(), SIMULATION_GIGAFLOW, 32); + return run_case("nuplan-gigaflow", drive_nuplan_map(), SIMULATION_MODE_GIGAFLOW, 32); } static int test_nuplan_replay_load_step_log(void) { - return run_case("nuplan-replay", drive_nuplan_map(), SIMULATION_REPLAY, 1); + return run_case("nuplan-replay", drive_nuplan_map(), SIMULATION_MODE_REPLAY, 1); } static int test_truncation_and_completed_episode_queue(void) { srand(11); - Drive env = drive_test_env_config(drive_carla_map(), SIMULATION_GIGAFLOW, 8, 0); + Drive env = drive_test_env_config(drive_carla_map(), SIMULATION_MODE_GIGAFLOW, 8, 0); env.scenario_length = 3; allocate(&env); c_reset(&env); diff --git a/tests/drive/test_drive_map_cache.c b/tests/drive/test_drive_map_cache.c index ef99a89b8a..ebb45417a1 100644 --- a/tests/drive/test_drive_map_cache.c +++ b/tests/drive/test_drive_map_cache.c @@ -7,7 +7,7 @@ static int test_obs_reward_done_parity_cache_on_vs_off(void) { const int steps = 20; srand(12345); - Drive off = drive_test_make_env(drive_carla_map(), SIMULATION_GIGAFLOW, 32, 0); + Drive off = drive_test_make_env(drive_carla_map(), SIMULATION_MODE_GIGAFLOW, 32, 0); int obs_count = off.active_agent_count * compute_observation_size(&off); int agent_count = off.active_agent_count; float *obs_log = (float *) malloc(steps * obs_count * sizeof(float)); @@ -26,7 +26,7 @@ static int test_obs_reward_done_parity_cache_on_vs_off(void) { free_allocated(&off); srand(12345); - Drive on = drive_test_make_env(drive_carla_map(), SIMULATION_GIGAFLOW, 32, 1); + Drive on = drive_test_make_env(drive_carla_map(), SIMULATION_MODE_GIGAFLOW, 32, 1); EXPECT_EQ_INT(on.active_agent_count, agent_count); EXPECT_EQ_INT(on.active_agent_count * compute_observation_size(&on), obs_count); for (int t = 0; t < steps; t++) { @@ -51,7 +51,7 @@ static int close_order_case(const int *order) { srand(9); Drive envs[3]; for (int i = 0; i < 3; i++) { - envs[i] = drive_test_make_env(drive_carla_map(), SIMULATION_GIGAFLOW, 8, 1); + envs[i] = drive_test_make_env(drive_carla_map(), SIMULATION_MODE_GIGAFLOW, 8, 1); drive_set_neutral_actions(&envs[i]); c_step(&envs[i]); } @@ -81,7 +81,7 @@ static int test_multi_env_close_orderings_do_not_crash(void) { static int test_cache_size_bounded_by_unique_maps(void) { drive_map_cache_clear(); for (int cycle = 0; cycle < 3; cycle++) { - Drive env = drive_test_make_env(drive_carla_map(), SIMULATION_GIGAFLOW, 8, 1); + Drive env = drive_test_make_env(drive_carla_map(), SIMULATION_MODE_GIGAFLOW, 8, 1); free_allocated(&env); EXPECT_EQ_INT(g_map_cache_count, 1); EXPECT_EQ_INT(drive_map_cache_live_count(), 0); @@ -92,7 +92,7 @@ static int test_cache_size_bounded_by_unique_maps(void) { static int test_forked_child_can_build_and_free_its_own_entry(void) { drive_map_cache_clear(); - Drive warm = drive_test_make_env(drive_carla_map(), SIMULATION_GIGAFLOW, 8, 1); + Drive warm = drive_test_make_env(drive_carla_map(), SIMULATION_MODE_GIGAFLOW, 8, 1); free_allocated(&warm); int parent_size_before_fork = g_map_cache_count; int fds[2]; @@ -101,7 +101,7 @@ static int test_forked_child_can_build_and_free_its_own_entry(void) { pid_t pid = fork(); if (pid == 0) { close(fds[0]); - Drive child = drive_test_make_env(drive_carla_map(), SIMULATION_GIGAFLOW, 8, 1); + Drive child = drive_test_make_env(drive_carla_map(), SIMULATION_MODE_GIGAFLOW, 8, 1); int live_after_build = drive_map_cache_live_count(); free_allocated(&child); int live_after_close = drive_map_cache_live_count(); diff --git a/tests/drive/test_drive_metrics.c b/tests/drive/test_drive_metrics.c index 10068c2ccb..2b6ba6f8b6 100644 --- a/tests/drive/test_drive_metrics.c +++ b/tests/drive/test_drive_metrics.c @@ -3,7 +3,7 @@ static int test_metric_offroad_outside_grid(void) { srand(5); - Drive env = drive_test_make_env(drive_carla_map(), SIMULATION_GIGAFLOW, 2, 0); + Drive env = drive_test_make_env(drive_carla_map(), SIMULATION_MODE_GIGAFLOW, 2, 0); int agent_idx = env.active_agent_indices[0]; Agent *agent = &env.agents[agent_idx]; agent->sim_x = env.grid_map->top_left_x - 1000.0f; @@ -40,7 +40,7 @@ static int test_metric_invalid_position_resets(void) { static int test_metric_on_road_lane_alignment(void) { srand(7); - Drive env = drive_test_make_env(drive_carla_map(), SIMULATION_GIGAFLOW, 1, 0); + Drive env = drive_test_make_env(drive_carla_map(), SIMULATION_MODE_GIGAFLOW, 1, 0); int agent_idx = env.active_agent_indices[0]; Agent *agent = &env.agents[agent_idx]; diff --git a/tests/drive/test_drive_observations_rewards.c b/tests/drive/test_drive_observations_rewards.c index d1b5a29a4c..4b6b80488f 100644 --- a/tests/drive/test_drive_observations_rewards.c +++ b/tests/drive/test_drive_observations_rewards.c @@ -3,7 +3,7 @@ static int test_observation_size_formula(void) { Drive env = {0}; - env.target_type = TARGET_STATIC; + env.target_type = TARGET_TYPE_STATIC; env.num_goals = 3; env.reward_conditioning = 0; env.obs_slots_partners_n = 2; @@ -14,7 +14,7 @@ static int test_observation_size_formula(void) { + 4 * TRAFFIC_CONTROL_FEATURES + OBS_VALID_COUNT_FEATURES; EXPECT_EQ_INT(compute_observation_size(&env), expected); - env.target_type = TARGET_DYNAMIC; + env.target_type = TARGET_TYPE_DYNAMIC; env.reward_conditioning = 1; expected = EGO_FEATURES + NUM_REWARD_COEFS + 3 * DYNAMIC_TARGET_FEATURES + 2 * PARTNER_FEATURES + 12 * ROAD_FEATURES + 4 * TRAFFIC_CONTROL_FEATURES + OBS_VALID_COUNT_FEATURES; @@ -24,7 +24,7 @@ static int test_observation_size_formula(void) { static int test_observation_zero_fill_and_valid_counts(void) { srand(3); - Drive env = drive_test_env_config(drive_carla_map(), SIMULATION_GIGAFLOW, 1, 0); + Drive env = drive_test_env_config(drive_carla_map(), SIMULATION_MODE_GIGAFLOW, 1, 0); env.obs_slots_partners_n = 4; env.obs_slots_lane_n = 8; env.obs_slots_boundary_n = 8; @@ -71,7 +71,7 @@ static void init_reward_env(Drive *env, Agent *agent, Log *log, int *active, flo env->active_agent_count = 1; env->dt = 0.1f; env->reward_goal = 2.0f; - env->simulation_mode = SIMULATION_GIGAFLOW; + env->simulation_mode = SIMULATION_MODE_GIGAFLOW; env->num_goals = 3; env->compute_eval_metrics = 0; agent->reward_coefs[REWARD_COEF_COLLISION] = 3.0f; diff --git a/tests/drive/test_perf.c b/tests/drive/test_perf.c index 2bf60a4a9f..247de67253 100644 --- a/tests/drive/test_perf.c +++ b/tests/drive/test_perf.c @@ -14,7 +14,7 @@ static int test_simulator_raw_perf(void) { const int baseline_sps = 24690; const float threshold = 0.8f * (float) baseline_sps; srand(17); - Drive env = drive_test_make_env(drive_carla_map(), SIMULATION_GIGAFLOW, 32, 0); + Drive env = drive_test_make_env(drive_carla_map(), SIMULATION_MODE_GIGAFLOW, 32, 0); EXPECT_EQ_INT(env.active_agent_count, 32); int tick = 0; diff --git a/tests/drive/test_render_smoke.c b/tests/drive/test_render_smoke.c index f5993087ea..47c034a828 100644 --- a/tests/drive/test_render_smoke.c +++ b/tests/drive/test_render_smoke.c @@ -4,7 +4,7 @@ static int test_render_default_frame(void) { srand(13); SetTraceLogLevel(LOG_WARNING); - Drive env = drive_test_env_config(drive_carla_map(), SIMULATION_GIGAFLOW, 8, 0); + Drive env = drive_test_env_config(drive_carla_map(), SIMULATION_MODE_GIGAFLOW, 8, 0); env.render_mode = RENDER_WINDOW; allocate(&env); c_reset(&env); diff --git a/tests/unit_tests/test_config_schema.py b/tests/unit_tests/test_config_schema.py new file mode 100644 index 0000000000..79825105e6 --- /dev/null +++ b/tests/unit_tests/test_config_schema.py @@ -0,0 +1,124 @@ +#!/usr/bin/env python3 +"""Structured-config schema tests: load_config validates the env section of +puffer_drive against pufferlib.config_schema.DriveEnvConfig at load time. + +Run: python -m unittest tests.unit_tests.test_config_schema +""" + +import os +import re +import sys +import unittest +from unittest.mock import patch + +from omegaconf.errors import ConfigKeyError, ValidationError + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) + +from pufferlib.config_schema import ( + ActionType, + ControlMode, + Controller, + DynamicsModel, + InfractionBehavior, + InitMode, + NonVehicleController, + SimulationMode, + TargetType, +) +from pufferlib.ocean.drive import binding +from pufferlib.pufferl import load_config + + +def _screaming_snake(name): + return re.sub(r"(?_`. This walks every enum + class instead of hand-picking members, so an enum added to + config_schema.py without a matching binding constant — or with a + mismatched value — fails here without needing a new assert line.""" + for enum_cls, member_suffix in _DRIFT_CHECKED_ENUMS: + group = _screaming_snake(enum_cls.__name__) + for member in enum_cls: + const_name = f"{group}_{member_suffix(member).upper()}" + self.assertTrue( + hasattr(binding, const_name), + f"binding.{const_name} not found for {enum_cls.__name__}.{member.name} " + "-- was the C #define renamed without updating config_schema.py, " + "or is it missing from env_binding.h's PyModule_AddIntConstant calls?", + ) + self.assertEqual( + getattr(binding, const_name), + member.value, + f"binding.{const_name} != {enum_cls.__name__}.{member.name}.value", + ) + + def test_non_vehicle_controller_matches_controller_constants(self): + """NonVehicleController reuses Controller's C constants; only its + config-only 'auto' sentinel (-1) has no C counterpart (drive.py + resolves it to a concrete Controller before it reaches C).""" + for member in NonVehicleController: + if member.name == "auto": + continue + const_name = f"CONTROLLER_{member.name.upper()}" + self.assertEqual(getattr(binding, const_name), member.value) + + +if __name__ == "__main__": + unittest.main()