diff --git a/.dockerignore b/.dockerignore index a896e4ff2d..e5abb4fde6 100644 --- a/.dockerignore +++ b/.dockerignore @@ -12,8 +12,6 @@ wandb .ruff_cache benchmark*/ -!pufferlib/ocean/benchmark/ -!pufferlib/ocean/benchmark/** runs*/ weights/ checkpoints/ diff --git a/.gitignore b/.gitignore index ee7286ba4a..aad05bd7c5 100644 --- a/.gitignore +++ b/.gitignore @@ -148,10 +148,6 @@ dmypy.json checkpoints/ experiments/ benchmark*/ -!pufferlib/ocean/benchmark/ -!pufferlib/ocean/benchmark/** -# But re-ignore caches inside it -pufferlib/ocean/benchmark/**/__pycache__/ wandb/ .neptune/ raylib*/ @@ -221,3 +217,7 @@ docs/_build/ # Claude config .claude/ CLAUDE.local.md + +episode_metrics + +/obs diff --git a/notebooks/05_inference.py b/notebooks/05_inference.py index ce8cfa3e80..5ad56d7bc6 100644 --- a/notebooks/05_inference.py +++ b/notebooks/05_inference.py @@ -24,7 +24,8 @@ from pufferlib.ocean.drive.drive import Drive from pufferlib.ocean.drive import binding from pufferlib.ocean.torch import Drive as DrivePolicy -from pufferlib.pytorch import sample_logits +from pufferlib.pytorch import sample_logits, ACTION_SELECT_MODE, ACTION_SELECT_SAMPLE +import pufferlib.spaces from notebooks.notebook_utils import COEF_NAMES, EGO_LABELS, MAP_DIR, load_notebook_config, zero_actions CHECKPOINT_PATH = "" @@ -55,8 +56,9 @@ policy.load_state_dict(sd) print(f"Loaded checkpoint: {CHECKPOINT_PATH}") -is_continuous = policy.is_continuous -ACT_SHAPE = (N, len(env.single_action_space.nvec)) if not is_continuous else (N, env.single_action_space.shape[0]) +is_continuous = policy.is_continuous # policy output space +env_continuous = isinstance(env.single_action_space, pufferlib.spaces.Box) # env action space +ACT_SHAPE = (N, env.single_action_space.shape[0]) if env_continuous else (N, len(env.single_action_space.nvec)) print(f"Policy on {device}, params: {sum(p.numel() for p in policy.parameters()):,}") print(f"Obs shape: {obs.shape}, Action space: {env.single_action_space}") @@ -77,8 +79,10 @@ logits_list, value = policy(obs_tensor) # Sample actions -action, logprob, ent = sample_logits(logits_list) -action_det, _, _ = sample_logits(logits_list, deterministic=True) +action, logprob, ent, cont_action = sample_logits(logits_list, env_continuous=env_continuous, policy=policy) +action_det, _, _, _ = sample_logits( + logits_list, action_selection=ACTION_SELECT_MODE, env_continuous=env_continuous, policy=policy +) print(f"Value: mean={value.mean():.4f}, std={value.std():.4f}, range=[{value.min():.4f}, {value.max():.4f}]") print(f"Entropy: mean={ent.mean():.4f}, std={ent.std():.4f}") @@ -122,7 +126,7 @@ n_tgt_wp = config["env"].get("num_goals", 3) -def run_rollout(env, policy, deterministic=False, horizon=HORIZON): +def run_rollout(env, policy, action_selection=ACTION_SELECT_SAMPLE, horizon=HORIZON): obs, _ = env.reset(seed=42) N = env.num_agents @@ -144,7 +148,9 @@ def run_rollout(env, policy, deterministic=False, horizon=HORIZON): obs_t = torch.FloatTensor(obs).to(device) with torch.no_grad(): logits_list, val = policy(obs_t) - act, logp, entr = sample_logits(logits_list, deterministic=deterministic) + act, logp, entr, cont_act = sample_logits( + logits_list, action_selection=action_selection, env_continuous=env_continuous, policy=policy + ) buffers["obs"][t] = obs buffers["actions"][t] = act.cpu().numpy().reshape(N) if act.dim() > 1 else act.cpu().numpy() @@ -157,8 +163,12 @@ def run_rollout(env, policy, deterministic=False, horizon=HORIZON): buffers["positions_x"][t] = gstate["x"] buffers["positions_y"][t] = gstate["y"] - # Step env - env_actions = act.cpu().numpy().reshape(ACT_SHAPE) + # Step env. A discrete policy on a continuous env steps with the decoded + # continuous action; otherwise the (discrete or continuous) action itself. + if env_continuous and not is_continuous: + env_actions = cont_act.cpu().numpy().reshape(N, -1) + else: + env_actions = act.cpu().numpy().reshape(ACT_SHAPE) obs, rew, term, trunc, info = env.step(env_actions) buffers["rewards"][t] = rew buffers["terminals"][t] = term @@ -168,9 +178,9 @@ def run_rollout(env, policy, deterministic=False, horizon=HORIZON): print("Running stochastic rollout...") -buf_stoch = run_rollout(env, policy, deterministic=False) +buf_stoch = run_rollout(env, policy, action_selection=ACTION_SELECT_SAMPLE) print("Running deterministic rollout...") -buf_det = run_rollout(env, policy, deterministic=True) +buf_det = run_rollout(env, policy, action_selection=ACTION_SELECT_MODE) for name, buf in [("Stochastic", buf_stoch), ("Deterministic", buf_det)]: print(f"\n--- {name} ---") @@ -1044,7 +1054,7 @@ def unpack_all_timesteps(bufs, agent_idx): # %% # Compute action probs over time for tracked agent (stochastic rollout) -n_actions = env.single_action_space.nvec[0] if not is_continuous else 1 +n_actions = policy.atn_dim[0] if not is_continuous else 1 action_probs_time = np.zeros((HORIZON, n_actions)) for t in range(HORIZON): obs_t = torch.FloatTensor(buf_stoch["obs"][t : t + 1, TRACKED_AGENT : TRACKED_AGENT + 1][0]).to(device) diff --git a/notebooks/06_architecture.py b/notebooks/06_architecture.py index 17e25077d6..671ff33c33 100644 --- a/notebooks/06_architecture.py +++ b/notebooks/06_architecture.py @@ -40,6 +40,7 @@ BACKBONE_ACTIVATION = "gelu" BACKBONE_LAYER_NORM = False MASK_PADDED_FEATURES = False +ACTION_TYPE = "discrete" env, obs, info = make_drive_env() @@ -64,6 +65,7 @@ backbone_layer_norm=BACKBONE_LAYER_NORM, shared_network=SHARED_NETWORK, mask_padded_features=MASK_PADDED_FEATURES, + action_type=ACTION_TYPE, ).to(device) print(f"Device: {device}") @@ -519,6 +521,7 @@ def count_params(module): "backbone_layer_norm": False, "shared_network": True, "mask_padded_features": False, + "action_type": "discrete", } results = [] diff --git a/notebooks/notebook_utils.py b/notebooks/notebook_utils.py index 41217c9029..3775d64388 100644 --- a/notebooks/notebook_utils.py +++ b/notebooks/notebook_utils.py @@ -112,6 +112,7 @@ "backbone_layer_norm": False, "shared_network": True, "mask_padded_features": False, + "action_type": "discrete", } diff --git a/pufferlib/config/puffer_drive.yaml b/pufferlib/config/puffer_drive.yaml index fec45d6bdb..a1b14fda1a 100644 --- a/pufferlib/config/puffer_drive.yaml +++ b/pufferlib/config/puffer_drive.yaml @@ -43,7 +43,7 @@ env: min_agents_per_env: 1 max_agents_per_env: 120 # Actions type - options: "discrete", "continuous" - action_type: discrete + action_type: continuous # Dynamics model - options: "classic", "jerk" dynamics_model: jerk # Time delta between steps in seconds @@ -191,6 +191,8 @@ policy: critic_num_layers: 1 # Dual or shared actor-critic backbone shared_network: false + action_type: discrete + rnn: input_size: 512 hidden_size: 512 @@ -280,6 +282,10 @@ eval: capture_observations: false observation_replay_wave_size: 16 observation_replay_writer_count: 4 + # Policy action selection during eval: sample | mode | mean. + # mode = argmax; mean = probability-weighted continuous action + # (mean requires a discrete policy on a continuous env). + action_selection: mode controlled_exp: train: # [sweep.train.learning_rate] diff --git a/pufferlib/ocean/drive/constants.h b/pufferlib/ocean/drive/constants.h index 4b31cddb71..25b17f19bc 100644 --- a/pufferlib/ocean/drive/constants.h +++ b/pufferlib/ocean/drive/constants.h @@ -137,12 +137,23 @@ static const float ACCEL_LONG_LIMIT[2] = {-5.0f, 2.5f}; static const float ACCEL_LAT_LIMIT[2] = {-4.0f, 4.0f}; // Discrete action space, DYNAMICS_MODEL_JERK -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}; +// These are the values used in the action discretization of the GIGAFLOW models. Potentially we might need to tune them +// for real car deployments. The limits should be set such that the model is forced to drive comfortably while still +// being able to control the car sufficiently as to drive well. The MIN limit for longetudinal (break) is much lower, +// such that the car can execute uncomfortable breaking behaviors in order to avoid collisions. Jerk action space (for +// JERK dynamics model) +#define NUM_JERK_LONG_ACTIONS 4 +#define NUM_JERK_LAT_ACTIONS 3 +static const float JERK_LONG[NUM_JERK_LONG_ACTIONS] = {-15.0f, -4.0f, 0.0f, 4.0f}; +static const float JERK_LAT[NUM_JERK_LAT_ACTIONS] = {-4.0f, 0.0f, 4.0f}; // Discrete action space, DYNAMICS_MODEL_CLASSIC -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}; +#define NUM_ACCELERATION_ACTIONS 7 +#define NUM_STEERING_ACTIONS 9 +static const float ACCELERATION_VALUES[NUM_ACCELERATION_ACTIONS] + = {-4.0000f, -2.6670f, -1.3330f, -0.0000f, 1.3330f, 2.6670f, 4.0000f}; +static const float STEERING_VALUES[NUM_STEERING_ACTIONS] + = {-0.667f, -0.500f, -0.333f, -0.167f, 0.000f, 0.167f, 0.333f, 0.500f, 0.667f}; // ===================================================================================== // 4. GEOMETRY & COLLISION diff --git a/pufferlib/ocean/drive/drive.h b/pufferlib/ocean/drive/drive.h index af66c913bb..55e239c50e 100644 --- a/pufferlib/ocean/drive/drive.h +++ b/pufferlib/ocean/drive/drive.h @@ -4121,7 +4121,6 @@ static void move_dynamics(Drive *env, int action_idx, int agent_idx) { j_long = JERK_LONG[0]; // max braking jerk j_lat = 0.0f; } - // Get dynamic conditioning coefficients float c_throttle = agent->reward_coefs[REWARD_COEF_THROTTLE]; float c_steer = agent->reward_coefs[REWARD_COEF_STEER]; diff --git a/pufferlib/ocean/env_binding.h b/pufferlib/ocean/env_binding.h index a0a19c1945..907a46754c 100644 --- a/pufferlib/ocean/env_binding.h +++ b/pufferlib/ocean/env_binding.h @@ -1313,6 +1313,28 @@ static PyMethodDef methods[] // Module definition static PyModuleDef module = {PyModuleDef_HEAD_INIT, "binding", NULL, -1, methods}; +// Publish a C float table as an immutable Python tuple so policy code reads the +// sim's values instead of copying them. Returns 0 on success, -1 on failure. +static int add_float_table_constant(PyObject *module_obj, const char *name, const float *values, int value_count) { + PyObject *table = PyTuple_New(value_count); + if (table == NULL) { + return -1; + } + for (int i = 0; i < value_count; i++) { + PyObject *value = PyFloat_FromDouble((double) values[i]); + if (value == NULL) { + Py_DECREF(table); + return -1; + } + PyTuple_SET_ITEM(table, i, value); + } + if (PyModule_AddObject(module_obj, name, table) < 0) { + Py_DECREF(table); + return -1; + } + return 0; +} + PyMODINIT_FUNC PyInit_binding(void) { import_array(); PyObject *m = PyModule_Create(&module); // Changed variable name from 'module' to 'm' @@ -1375,5 +1397,15 @@ PyMODINIT_FUNC PyInit_binding(void) { 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)); + // Action discretization tables: the policy decodes discrete actions with the + // exact values the sim uses, so the two can never drift apart. + if (add_float_table_constant(m, "JERK_LONG", JERK_LONG, NUM_JERK_LONG_ACTIONS) < 0 + || add_float_table_constant(m, "JERK_LAT", JERK_LAT, NUM_JERK_LAT_ACTIONS) < 0 + || add_float_table_constant(m, "ACCELERATION_VALUES", ACCELERATION_VALUES, NUM_ACCELERATION_ACTIONS) < 0 + || add_float_table_constant(m, "STEERING_VALUES", STEERING_VALUES, NUM_STEERING_ACTIONS) < 0) { + Py_DECREF(m); + return NULL; + } + return m; } diff --git a/pufferlib/ocean/evaluation_utils/wosac/evaluator.py b/pufferlib/ocean/evaluation_utils/wosac/evaluator.py index 2798e0a926..a325fcfc7f 100644 --- a/pufferlib/ocean/evaluation_utils/wosac/evaluator.py +++ b/pufferlib/ocean/evaluation_utils/wosac/evaluator.py @@ -88,6 +88,8 @@ def collect_simulated_trajectories(self, args, puffer_env, policy): "id": np.zeros((num_agents, self.num_rollouts, self.sim_steps), dtype=np.int32), } + env_continuous = isinstance(puffer_env.single_action_space, pufferlib.spaces.Box) + for rollout_idx in range(self.num_rollouts): print(f"\rCollecting rollout {rollout_idx + 1}/{self.num_rollouts}...", end="", flush=True) obs, info = puffer_env.reset() @@ -98,6 +100,7 @@ def collect_simulated_trajectories(self, args, puffer_env, policy): lstm_c=torch.zeros(num_agents, policy.hidden_size, device=device), ) + # TODO find out how to test this code path for time_idx in range(self.sim_steps): # Get global state agent_state = driver.get_global_agent_state() @@ -111,13 +114,19 @@ def collect_simulated_trajectories(self, args, puffer_env, policy): with torch.no_grad(): ob_tensor = torch.as_tensor(obs).to(device) logits, value = policy.forward_eval(ob_tensor, state) - action, logprob, _ = pufferlib.pytorch.sample_logits(logits) + action, logprob, _, cont_action = pufferlib.pytorch.sample_logits( + logits, env_continuous=env_continuous, policy=policy + ) action_np = action.cpu().numpy().reshape(puffer_env.action_space.shape) if isinstance(logits, torch.distributions.Normal): action_np = np.clip(action_np, puffer_env.action_space.low, puffer_env.action_space.high) - obs, _, _, _, _ = puffer_env.step(action_np) + if env_continuous and not policy.is_continuous: + cont_action = cont_action.cpu().numpy().reshape(puffer_env.action_space.shape) + obs, _, _, _, _ = puffer_env.step(cont_action) + else: + obs, _, _, _, _ = puffer_env.step(action_np) return trajectories @@ -659,6 +668,7 @@ def rollout(self, args, puffer_env, policy): total_steps = (scenario_length - init_steps + 1) * num_maps obs, info = puffer_env.reset() + env_continuous = isinstance(puffer_env.single_action_space, pufferlib.spaces.Box) state = {} if args["train"]["use_rnn"]: state = dict( @@ -671,13 +681,20 @@ def rollout(self, args, puffer_env, policy): with torch.no_grad(): ob_tensor = torch.as_tensor(obs).to(device) logits, value = policy.forward_eval(ob_tensor, state) - action, logprob, _ = pufferlib.pytorch.sample_logits(logits) + action, logprob, _, cont_action = pufferlib.pytorch.sample_logits( + logits, env_continuous=env_continuous, policy=policy + ) action_np = action.cpu().numpy().reshape(puffer_env.action_space.shape) if isinstance(logits, torch.distributions.Normal): action_np = np.clip(action_np, puffer_env.action_space.low, puffer_env.action_space.high) - obs, rewards, dones, truncs, info_list = puffer_env.step(action_np) + if env_continuous and not policy.is_continuous: + cont_action = cont_action.cpu().numpy().reshape(puffer_env.action_space.shape) + obs, rewards, dones, truncs, info_list = puffer_env.step(cont_action) + else: + obs, rewards, dones, truncs, info_list = puffer_env.step(action_np) + if info_list: all_infos.extend(info_list) # Stop once we've collected at least one info per bin to avoid diff --git a/pufferlib/ocean/torch.py b/pufferlib/ocean/torch.py index 858278f45c..fb703d78e2 100644 --- a/pufferlib/ocean/torch.py +++ b/pufferlib/ocean/torch.py @@ -330,9 +330,43 @@ def __init__( backbone_layer_norm: bool, shared_network: bool, mask_padded_features: bool, + action_type: str, ): super().__init__() + # Action discretization tables come from the sim (drive.h, exposed by the + # C binding) so the policy's decoding can never drift from the env's. + if env.dynamics_model == "jerk": + action_long_values, action_lat_values = binding.JERK_LONG, binding.JERK_LAT + elif env.dynamics_model == "classic": + action_long_values, action_lat_values = binding.ACCELERATION_VALUES, binding.STEERING_VALUES + else: + raise ValueError(f"Unsupported dynamics model: {env.dynamics_model}") + + action_long = torch.tensor(action_long_values, dtype=torch.float32) + action_lat = torch.tensor(action_lat_values, dtype=torch.float32) + + # Precompute the [-1, 1] continuous action per discrete choice by inverting the + # sim's continuous scaling (drive.h). Constant → done once. Longitudinal is + # asymmetric (braking / |t[0]|, accel / t[-1]); lateral symmetric. The symmetric + # classic table collapses both branches to / t[-1]. + long_norm = torch.where( + action_long < 0.0, + action_long / -action_long[0], + action_long / action_long[-1], + ) + lat_norm = action_lat / action_lat[-1] + self.register_buffer("action_long_norm", long_norm, persistent=False) + self.register_buffer("action_lat_norm", lat_norm, persistent=False) + + # Joint continuous-action table: row k = the [-1,1] (long, lat) for discrete class k, + # where k = long_idx * num_lat + lat_idx (matches drive.h decode). + num_lat = lat_norm.numel() + num_classes = long_norm.numel() * num_lat + k = torch.arange(num_classes) + action_table = torch.stack([long_norm[k // num_lat], lat_norm[k % num_lat]], dim=-1) # [num_classes, 2] + self.register_buffer("action_table", action_table, persistent=False) + # Configuration flags from policy kwargs self.shared_network = shared_network self.ego_dim = env.ego_features @@ -367,11 +401,11 @@ def __init__( self.critic_backbone = DriveBackbone(**backbone_args) # Setup action and value heads - self.is_continuous = isinstance(env.single_action_space, pufferlib.spaces.Box) + self.is_continuous = action_type == "continuous" if self.is_continuous: self.atn_dim = (env.single_action_space.shape[0],) * 2 else: - self.atn_dim = env.single_action_space.nvec.tolist() + self.atn_dim = [self.action_long_norm.numel() * self.action_lat_norm.numel()] # n-layer MLP for actor head (num_layers = number of hidden layers) backbone_out_dim = self.actor_backbone.out_dim @@ -454,3 +488,10 @@ def decode_actions(self, hidden): value = self.critic_head(hidden) return action, value + + def discrete_actions_to_continuous(self, actions): + return self.action_table[actions.long()] + + def discrete_probs_to_continuous_mean(self, probs): + # probs: [..., num_classes] -> [..., 2] (E[cont | probs]) + return probs @ self.action_table.to(probs.dtype) diff --git a/pufferlib/pufferl.py b/pufferlib/pufferl.py index afc0c69e20..a230dee932 100644 --- a/pufferlib/pufferl.py +++ b/pufferlib/pufferl.py @@ -127,8 +127,17 @@ def __init__(self, config, vecenv, policy, logger=None): # Vecenv info vecenv.async_reset(seed) + + self.env_continuous = isinstance(vecenv.single_action_space, pufferlib.spaces.Box) obs_space = vecenv.single_observation_space - atn_space = vecenv.single_action_space + # Custom policy attributes live on the base module, not the DDP/compile wrapper. + unwrapped_policy = base_policy(policy) + if self.env_continuous and not unwrapped_policy.is_continuous: + action_shape = (len(unwrapped_policy.atn_dim),) + action_dtype = torch.int32 + else: + action_shape = vecenv.single_action_space.shape + action_dtype = pufferlib.pytorch.numpy_to_torch_dtype_dict[vecenv.single_action_space.dtype] total_agents = vecenv.num_agents self.total_agents = total_agents @@ -170,9 +179,9 @@ def __init__(self, config, vecenv, policy, logger=None): self.actions = torch.zeros( segments, horizon, - *atn_space.shape, + *action_shape, device=device, - dtype=pufferlib.pytorch.numpy_to_torch_dtype_dict[atn_space.dtype], + dtype=action_dtype, ) self.values = torch.zeros(segments, horizon, device=device) self.logprobs = torch.zeros(segments, horizon, device=device) @@ -396,7 +405,9 @@ def evaluate(self): logits, value = self.policy.forward_eval(o_device, state) logits = logits_to_float(logits) - action, logprob, _ = pufferlib.pytorch.sample_logits(logits) + action, logprob, _, cont_action = pufferlib.pytorch.sample_logits( + logits, env_continuous=self.env_continuous, policy=self.uncompiled_policy + ) if config["normalize_rewards"]: r = torch.sign(r) * torch.log1p(torch.abs(r)) @@ -439,10 +450,6 @@ def evaluate(self): self.free_idx += num_full self.full_rows += num_full - action = action.cpu().numpy() - if isinstance(logits, torch.distributions.Normal): - action = np.clip(action, self.vecenv.action_space.low, self.vecenv.action_space.high) - profile("eval_misc", epoch) for i in info: for k, v in pufferlib.unroll_nested_dict(i): @@ -454,7 +461,15 @@ def evaluate(self): self.stats[k].append(v) profile("env", epoch) - self.vecenv.send(action) + + if self.env_continuous and not self.uncompiled_policy.is_continuous: + cont_action = cont_action.cpu().numpy() + self.vecenv.send(cont_action.squeeze(0)) + else: + action = action.cpu().numpy() + if isinstance(logits, torch.distributions.Normal): + action = np.clip(action, self.vecenv.action_space.low, self.vecenv.action_space.high) + self.vecenv.send(action) profile("eval_misc", epoch) self.free_idx = self.total_agents @@ -534,7 +549,7 @@ def _ppo_loss(self, mb_obs, mb_actions, mb_logprobs, mb_values, mb_returns, mb_a logits, newvalue = self.policy(mb_obs, state) logits = logits_to_float(logits) newvalue = newvalue.float() - _, newlogprob, entropy = pufferlib.pytorch.sample_logits(logits, action=mb_actions) + _, newlogprob, entropy, _ = pufferlib.pytorch.sample_logits(logits, action=mb_actions) newlogprob = newlogprob.float().view_as(mb_logprobs) newvalue = newvalue.view_as(mb_returns) @@ -546,6 +561,7 @@ def _ppo_loss(self, mb_obs, mb_actions, mb_logprobs, mb_values, mb_returns, mb_a approx_kl = ((ratio - 1) - logratio).mean() clipfrac = ((ratio - 1.0).abs() > config["clip_coef"]).float().mean() + # TODO fix multi-gpu bug mb_adv = (mb_adv - mb_adv.mean()) / (mb_adv.std(unbiased=False) + 1e-8) if adv_weights is not None: mb_adv = adv_weights * mb_adv @@ -1580,6 +1596,15 @@ def eval( max_rendered_failures = eval_config["max_rendered_failures"] failure_replay_csv = eval_config["failure_replay_csv"] max_sdc_replay_workers = eval_config["max_sdc_replay_workers"] + valid_action_selections = ( + pufferlib.pytorch.ACTION_SELECT_SAMPLE, + pufferlib.pytorch.ACTION_SELECT_MODE, + pufferlib.pytorch.ACTION_SELECT_MEAN, + ) + if eval_config["action_selection"] not in valid_action_selections: + raise pufferlib.APIUsageError( + f"eval.action_selection='{eval_config['action_selection']}' must be one of {valid_action_selections}" + ) if render_scenarios and failure_replay_csv is not None: raise pufferlib.APIUsageError( "eval.render_scenarios requires a standard benchmark pass and cannot be combined " @@ -1928,6 +1953,12 @@ def _run_eval_rollout( evaluation_policy_cache["sample_logits"] = eval_sample_logits policy_forward_eval = evaluation_policy_cache["policy_forward_eval"] eval_sample_logits = evaluation_policy_cache["sample_logits"] + # A discrete policy on a continuous env emits a discrete class that the + # policy's own table maps back to the continuous action the env expects. + action_selection = args["eval"]["action_selection"] + uncompiled_policy = base_policy(policy) + env_continuous = isinstance(vecenv.single_action_space, pufferlib.spaces.Box) + discrete_policy_on_continuous_env = env_continuous and not uncompiled_policy.is_continuous device = torch_device(args["train"]["device"]) use_bfloat16 = args["train"]["amp"] and args["train"]["precision"] == "bfloat16" and is_cuda_device(device) if use_bfloat16 and not torch.cuda.is_bf16_supported(): @@ -1976,9 +2007,21 @@ def _run_eval_rollout( logits, value = policy_forward_eval(policy_obs_tensor) else: logits, value = policy_forward_eval(policy_obs_tensor, recurrent_state) - action, logprob, entropy = eval_sample_logits(logits, deterministic=True) - raw_action = action[:agents_per_batch].cpu().numpy().reshape(vecenv.action_space.shape) - action = raw_action + action, logprob, entropy, cont_action = eval_sample_logits( + logits, + action_selection=action_selection, + env_continuous=env_continuous, + policy=uncompiled_policy, + ) + if discrete_policy_on_continuous_env: + # raw_action stays the discrete class (what the replay logs record), + # while the env is stepped with its continuous counterpart. + raw_action = action[:agents_per_batch].cpu().numpy() + continuous_actions = cont_action.reshape(-1, *vecenv.single_action_space.shape) + action = continuous_actions[:agents_per_batch].float().cpu().numpy() + else: + raw_action = action[:agents_per_batch].cpu().numpy().reshape(vecenv.action_space.shape) + action = raw_action if isinstance(logits, torch.distributions.Normal): action = np.clip(action, vecenv.action_space.low, vecenv.action_space.high) diff --git a/pufferlib/pytorch.py b/pufferlib/pytorch.py index 2ec10a5c5f..96d2e39078 100644 --- a/pufferlib/pytorch.py +++ b/pufferlib/pytorch.py @@ -20,6 +20,9 @@ np.dtype("int8"): torch.int8, } +ACTION_SELECT_SAMPLE = "sample" +ACTION_SELECT_MODE = "mode" # argmax +ACTION_SELECT_MEAN = "mean" # probability-weighted continuous mean LITTLE_BYTE_ORDER = sys.byteorder == "little" @@ -187,18 +190,22 @@ def entropy_probs(logits, probs): return -p_log_p.sum(-1) -def sample_logits(logits, action=None, deterministic=False): +def sample_logits( + logits, action=None, action_selection=ACTION_SELECT_SAMPLE, env_continuous=None, policy=None +): # TODO discrete continuous is_discrete = isinstance(logits, torch.Tensor) + if action_selection == ACTION_SELECT_MEAN and not (env_continuous and not policy.is_continuous): + raise ValueError("action_selection='mean' requires a discrete policy on a continuous env") if isinstance(logits, torch.distributions.Normal): batch = logits.loc.shape[0] if action is None: action = logits.sample().view(batch, -1) - if deterministic: - action = logits.loc.view(batch, -1) # TODO - DETERMINISTIC use mean action for eval + if action_selection != ACTION_SELECT_SAMPLE: + action = logits.loc.view(batch, -1) log_probs = logits.log_prob(action.view(batch, -1)).sum(1) logits_entropy = logits.entropy().view(batch, -1).sum(1) - return action, log_probs, logits_entropy + return action, log_probs, logits_entropy, None elif is_discrete: logits = logits.unsqueeze(0) # TODO: Double check this @@ -212,12 +219,12 @@ def sample_logits(logits, action=None, deterministic=False): probs = logits_to_probs(logits) if action is None: - if deterministic: - action = torch.argmax(probs, -1) - else: + if action_selection == ACTION_SELECT_SAMPLE: probs = torch.nan_to_num(probs, 1e-8, 1e-8, 1e-8) action = torch.multinomial(probs.reshape(-1, probs.shape[-1]), 1, replacement=True).int() action = action.reshape(probs.shape[:-1]) + else: # MODE and MEAN both use argmax as the *nominal* discrete action (for logging) + action = torch.argmax(probs, -1) else: batch = logits[0].shape[0] action = action.view(batch, -1).T @@ -226,7 +233,16 @@ def sample_logits(logits, action=None, deterministic=False): logprob = log_prob(normalized_logits, action) logits_entropy = entropy(normalized_logits).sum(0) + if env_continuous and not policy.is_continuous: + if action_selection == ACTION_SELECT_MEAN: + cont_actions = policy.discrete_probs_to_continuous_mean(probs) + else: # Mode and sample we already have an action from the code above. + cont_actions = policy.discrete_actions_to_continuous(action) + if is_discrete: + return action.squeeze(0), logprob.squeeze(0), logits_entropy.squeeze(0), cont_actions.squeeze(0) + return action.T, logprob.sum(0), logits_entropy, cont_actions + if is_discrete: - return action.squeeze(0), logprob.squeeze(0), logits_entropy.squeeze(0) + return action.squeeze(0), logprob.squeeze(0), logits_entropy.squeeze(0), None - return action.T, logprob.sum(0), logits_entropy + return action.T, logprob.sum(0), logits_entropy, None diff --git a/scripts/baseline_10G/config-reward-dense.yaml b/scripts/baseline_10G/config-reward-dense.yaml index f9eed82808..77922f036d 100644 --- a/scripts/baseline_10G/config-reward-dense.yaml +++ b/scripts/baseline_10G/config-reward-dense.yaml @@ -324,6 +324,7 @@ policy: partner_input_size: 128 shared_network: true traffic_control_input_size: 128 + action_type: discrete policy_name: Drive render: 0 render_mode: auto diff --git a/scripts/baseline_10G/config-reward-sparse.yaml b/scripts/baseline_10G/config-reward-sparse.yaml index 645d7c9523..77497998a0 100644 --- a/scripts/baseline_10G/config-reward-sparse.yaml +++ b/scripts/baseline_10G/config-reward-sparse.yaml @@ -324,6 +324,7 @@ policy: partner_input_size: 128 shared_network: true traffic_control_input_size: 128 + action_type: discrete policy_name: Drive render: 0 render_mode: auto diff --git a/scripts/cluster_configs/nightly_best.yaml b/scripts/cluster_configs/nightly_best.yaml index 8693b9552f..0795108ed7 100644 --- a/scripts/cluster_configs/nightly_best.yaml +++ b/scripts/cluster_configs/nightly_best.yaml @@ -9,6 +9,7 @@ # Full training scale. The defaults are sized to run on-premise for dev; # the scale knobs for a real training run live here. +policy.action_type: discrete env.num_agents: 3200 vec.num_envs: 40 train.minibatch_size: 256000 diff --git a/scripts/cluster_configs/single_agent_speed_run.yaml b/scripts/cluster_configs/single_agent_speed_run.yaml index fa555ba2bd..b6dd076258 100644 --- a/scripts/cluster_configs/single_agent_speed_run.yaml +++ b/scripts/cluster_configs/single_agent_speed_run.yaml @@ -67,6 +67,7 @@ policy.backbone_num_layers: 2 policy.actor_num_layers: 0 policy.critic_num_layers: 0 policy.shared_network: true +policy.action_type: discrete # Training: short schedule, small minibatch, compiled train.total_timesteps: 1_000_000_000 diff --git a/scripts/config-reward-dense.yaml b/scripts/config-reward-dense.yaml new file mode 100644 index 0000000000..77922f036d --- /dev/null +++ b/scripts/config-reward-dense.yaml @@ -0,0 +1,403 @@ +agent_index: null +controlled_exp: + train: + ent_coef: + values: + - 0.01 + - 0.005 + learning_rate: + values: + - 0.001 + - 0.003 + - 0.01 +env: + action_type: discrete + collision_behavior: 1 + compute_eval_metrics: false + control_mode: control_vehicles + dt: 0.1 + dynamics_model: jerk + goal_on_lane: true + goal_radius: 2.0 + goal_speed: 1000.0 + inactive_agent_threshold: 0.4 + init_mode: create_all_valid + init_step: 0 + map_dir: "pufferlib/resources/drive/binaries/carla" + max_agents_per_env: 80 + max_waypoint_spacing: 40.0 + min_agents_per_env: 1 + min_waypoint_spacing: 20.0 + non_sdc_controller: policy + non_vehicle_controller: auto + num_agents: 2048 + num_maps: 8 + num_target_waypoints: 3 + obs_dropout_boundary: 0.0 + obs_dropout_lane: 0.0 + obs_norm_goal_offset_m: 120.0 + obs_norm_road_seg_length_m: 10.0 + obs_norm_road_seg_width_m: 5.0 + obs_norm_veh_length_m: 10.0 + obs_norm_veh_width_m: 5.0 + obs_norm_xy_offset_m: 120.0 + obs_range_partner_m: 120.0 + obs_range_road_behind_m: 30.0 + obs_range_road_front_m: 120.0 + obs_range_road_side_m: 40.0 + obs_range_traffic_control_m: 120.0 + obs_slots_boundary_n: 30 + obs_slots_lane_n: 50 + obs_slots_partners_n: 12 + obs_slots_traffic_controls_n: 4 + offroad_behavior: 1 + partner_blindness_prob: 0.0 + partner_blindness_trigger_prob: 0.0 + phantom_braking_duration: 10 + phantom_braking_prob: 0.0 + phantom_braking_trigger_prob: 0.0 + resample_frequency: 102400 + reward_ade: 0.0 + reward_center_bias: 0.0 + reward_collision: 1.5 + reward_comfort: 0.5 + reward_conditioning: false + reward_goal: 0.2 + reward_lane_align: 0.5 + reward_lane_center: 0.05 + reward_offroad: 1.5 + reward_overspeed: 0.1 + reward_randomization: false + reward_reverse: 0.005 + reward_stop_line: 1.0 + reward_timestep: 0.0 + reward_vel_align: 1.0 + reward_velocity: 0.1 + scenario_length: 1024 + sdc_controller: policy + simulation_mode: gigaflow + spawn_initial_speed: 0.0 + goal_regen_mode: finite + goal_source: route + termination_mode: 1 + traffic_light_behavior: 1 + use_map_cache: 1 +env_name: puffer_drive +eval: + behaviors_defaults: + clean: 'true' + enabled: 'false' + env: + control_mode: control_sdc_only + init_mode: create_all_valid + obs_slots_partners_n: 32 + scenario_length: 201 + simulation_mode: replay + eval: + num_scenarios: 50 + render_max_steps: 200 + render_num_scenarios: 2 + interval: 250 + mode: inline + render: 'true' + render_views: + - sim_state + - bev + behaviors_full_dir: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/nuplan_mini_train_bins + inherits: behaviors_defaults + type: behavior_class + behaviors_hard_stop: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/categories_v021/hard_stop + inherits: behaviors_defaults + type: behavior_class + behaviors_highway_straight: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/categories_v021/highway_straight + inherits: behaviors_defaults + type: behavior_class + behaviors_lane_change: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/categories_v021/lane_change + inherits: behaviors_defaults + type: behavior_class + behaviors_merge: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/categories_v021/merge + inherits: behaviors_defaults + type: behavior_class + behaviors_parked_cars: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/categories_v021/parked_cars + inherits: behaviors_defaults + type: behavior_class + behaviors_roundabout: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/categories_v021/roundabout + inherits: behaviors_defaults + type: behavior_class + behaviors_stopped_traffic: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/categories_v021/stopped_traffic + inherits: behaviors_defaults + type: behavior_class + behaviors_traffic_light_green: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/categories_v021/traffic_light_green + inherits: behaviors_defaults + type: behavior_class + behaviors_traffic_light_stop: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/categories_v021/traffic_light_stop + inherits: behaviors_defaults + type: behavior_class + behaviors_unprotected_left: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/categories_v021/unprotected_left + inherits: behaviors_defaults + type: behavior_class + behaviors_unprotected_right: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/categories_v021/unprotected_right + inherits: behaviors_defaults + type: behavior_class + dnf_triage: + enabled: 'false' + env: + map_dir: pufferlib/resources/drive/binaries/carla/opendrive__Town10HD.bin + max_agents_per_env: 1 + min_agents_per_env: 1 + num_maps: 1 + resample_frequency: 500 + scenario_length: 500 + simulation_mode: gigaflow + eval: + num_scenarios: 32 + render_max_steps: 300 + render_num_scenarios: 16 + inherits: validation_defaults + render: 'true' + render_backend: triage_html + type: multi_scenario + validation_defaults: + clean: 'true' + enabled: 'true' + env: + collision_behavior: 1 + eval_mode: 1 + goal_speed: 3.0 + num_agents: 1024 + obs_dropout_boundary: 0.0 + obs_dropout_lane: 0.0 + obs_slots_boundary_n: 80 + obs_slots_lane_n: 80 + offroad_behavior: 1 + reward_ade: 0.0 + reward_collision: 3.0 + reward_comfort: 0.05 + reward_goal: 1.0 + reward_lane_align: 0.025 + reward_lane_center: 0.0038 + reward_offroad: 3.0 + reward_overspeed: 0.05 + reward_randomization: false + reward_reverse: 0.005 + reward_stop_line: 1.0 + reward_timestep: 2.5e-05 + reward_velocity: 0.0025 + goal_source: route + goal_regen_mode: finite + termination_mode: 0 + traffic_light_behavior: 0 + eval: + export_episode_csv: 'true' + num_scenarios: 250 + verify_coverage: 'true' + interval: 250 + mode: inline + validation_gigaflow: + enabled: 'true' + env: + map_dir: pufferlib/resources/drive/binaries/carla + max_agents_per_env: 40 + min_agents_per_env: 40 + num_agents: 1024 + num_maps: 8 + resample_frequency: 500 + scenario_length: 500 + simulation_mode: gigaflow + eval: + render_max_steps: 300 + render_num_scenarios: 8 + inherits: validation_defaults + render: 'true' + render_backend: egl + render_views: + - sim_state + - bev + type: multi_scenario + validation_replay: + enabled: 'true' + env: + control_mode: control_sdc_only + map_dir: /scratch/ev2237/data/nuplan/nuplan_mini_train_bins + max_agents_per_env: 64 + num_maps: 250 + resample_frequency: 200 + scenario_length: 200 + simulation_mode: replay + eval: + render_max_steps: 200 + render_num_scenarios: 5 + inherits: validation_defaults + render: 'true' + render_backend: triage_html + type: multi_scenario + wosac: + clean: 'true' + enabled: 'false' + env: + control_mode: control_wosac + goal_radius: 2.0 + init_mode: create_all_valid + init_step: 10 + eval: + wosac_aggregate_results: 'true' + wosac_num_agents: 256 + wosac_num_rollouts: 32 + wosac_sanity_check: 'false' + interval: 500 + mode: subprocess + render: 'false' + type: wosac +eval_simulation: null +fps: 15 +gif_path: eval.gif +git: + commit_hash: 9d3fcc5c09db50b9674a374b8c3404b6349c4cb5 +load_id: null +load_model_path: null +local_rank: 0 +max_runs: 200 +max_suggestion_cost: 3600 +mine: + num_episodes: 100 + output_dir: '' + render: 'true' + score_threshold: -inf +neptune: false +neptune_name: pufferai +neptune_project: ablations +no_model_upload: {} +num_scenarios: 3 +package: ocean +policy: + actor_hidden_size: 512 + actor_num_layers: 0 + backbone_activation: gelu + backbone_hidden_size: 512 + backbone_layer_norm: false + backbone_num_layers: 2 + boundary_input_size: 128 + context_input_size: 128 + critic_hidden_size: 512 + critic_num_layers: 0 + ego_input_size: 128 + encoder_activation: relu + encoder_layer_norm: true + lane_input_size: 128 + mask_padded_features: false + partner_input_size: 128 + shared_network: true + traffic_control_input_size: 128 + action_type: discrete +policy_name: Drive +render: 0 +render_mode: auto +rnn: + hidden_size: 512 + input_size: 512 +rnn_name: null +run_name: rew-dense-seed4 +save_frames: 0 +sweep: + downsample: 10 + goal: maximize + method: Protein + metric: score +tag: null +tb: false +train: + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_eps: 1.0e-08 + adv_filter_ewma_beta: 0.25 + adv_filter_threshold_scale: 0.01 + adv_sampling_prio_alpha: 0.8499999999999999 + adv_sampling_prio_beta0: 0.8499999999999999 + amp: true + anneal_lr: true + batch_size: auto + bptt_horizon: 128 + checkpoint_interval: 50 + clip_coef: 0.2 + compile: true + compile_fullgraph: false + compile_mode: default + cpu_offload: false + data_dir: /pufferdrive/training_output/ + device: cuda + ent_coef: 0.01 + gae_lambda: 0.95 + gamma: 0.999 + learning_rate: 0.0005 + max_grad_norm: 0.5 + max_minibatch_size: 98304 + minibatch_size: 98304 + name: pufferai + normalize_rewards: true + obs_only: true + optimizer: adamw + precision: bfloat16 + project: ablations + render: false + render_interval: 1000 + render_map: none + resume_state_path: null + seed: 4 + show_grid: false + show_human_logs: true + show_lasers: false + torch_deterministic: false + total_timesteps: 10000000000 + update_epochs: 2 + use_rnn: false + vf_clip_coef: null + vf_coef: 0.5 + vtrace_c_clip: 1 + vtrace_rho_clip: 1 +vec: + backend: Multiprocessing + batch_size: auto + num_envs: 24 + num_workers: auto + seed: 42 + zero_copy: true +video_path: videos +wandb: true +wandb_group: pr-smoke-runs +wandb_project: pr-smoke-runs diff --git a/scripts/config-reward-sparse.yaml b/scripts/config-reward-sparse.yaml new file mode 100644 index 0000000000..77497998a0 --- /dev/null +++ b/scripts/config-reward-sparse.yaml @@ -0,0 +1,403 @@ +agent_index: null +controlled_exp: + train: + ent_coef: + values: + - 0.01 + - 0.005 + learning_rate: + values: + - 0.001 + - 0.003 + - 0.01 +env: + action_type: discrete + collision_behavior: 1 + compute_eval_metrics: false + control_mode: control_vehicles + dt: 0.1 + dynamics_model: jerk + goal_on_lane: true + goal_radius: 2.0 + goal_speed: 1000.0 + inactive_agent_threshold: 0.4 + init_mode: create_all_valid + init_step: 0 + map_dir: "pufferlib/resources/drive/binaries/carla" + max_agents_per_env: 80 + max_waypoint_spacing: 40.0 + min_agents_per_env: 1 + min_waypoint_spacing: 20.0 + non_sdc_controller: policy + non_vehicle_controller: auto + num_agents: 2048 + num_maps: 8 + num_target_waypoints: 3 + obs_dropout_boundary: 0.0 + obs_dropout_lane: 0.0 + obs_norm_goal_offset_m: 120.0 + obs_norm_road_seg_length_m: 10.0 + obs_norm_road_seg_width_m: 5.0 + obs_norm_veh_length_m: 10.0 + obs_norm_veh_width_m: 5.0 + obs_norm_xy_offset_m: 120.0 + obs_range_partner_m: 120.0 + obs_range_road_behind_m: 30.0 + obs_range_road_front_m: 120.0 + obs_range_road_side_m: 40.0 + obs_range_traffic_control_m: 120.0 + obs_slots_boundary_n: 30 + obs_slots_lane_n: 50 + obs_slots_partners_n: 12 + obs_slots_traffic_controls_n: 4 + offroad_behavior: 1 + partner_blindness_prob: 0.0 + partner_blindness_trigger_prob: 0.0 + phantom_braking_duration: 10 + phantom_braking_prob: 0.0 + phantom_braking_trigger_prob: 0.0 + resample_frequency: 102400 + reward_ade: 0.0 + reward_center_bias: 0.0 + reward_collision: 1.5 + reward_comfort: 0.05 + reward_conditioning: false + reward_goal: 0.5 + reward_lane_align: 0.025 + reward_lane_center: 0.0038 + reward_offroad: 1.5 + reward_overspeed: 0.05 + reward_randomization: false + reward_reverse: 0.005 + reward_stop_line: 1.0 + reward_timestep: 0.0 + reward_vel_align: 1.0 + reward_velocity: 0.0025 + scenario_length: 1024 + sdc_controller: policy + simulation_mode: gigaflow + spawn_initial_speed: 0.0 + goal_source: route + goal_regen_mode: finite + termination_mode: 1 + traffic_light_behavior: 1 + use_map_cache: 1 +env_name: puffer_drive +eval: + behaviors_defaults: + clean: 'true' + enabled: 'false' + env: + control_mode: control_sdc_only + init_mode: create_all_valid + obs_slots_partners_n: 32 + scenario_length: 201 + simulation_mode: replay + eval: + num_scenarios: 50 + render_max_steps: 200 + render_num_scenarios: 2 + interval: 250 + mode: inline + render: 'true' + render_views: + - sim_state + - bev + behaviors_full_dir: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/nuplan_mini_train_bins + inherits: behaviors_defaults + type: behavior_class + behaviors_hard_stop: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/categories_v021/hard_stop + inherits: behaviors_defaults + type: behavior_class + behaviors_highway_straight: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/categories_v021/highway_straight + inherits: behaviors_defaults + type: behavior_class + behaviors_lane_change: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/categories_v021/lane_change + inherits: behaviors_defaults + type: behavior_class + behaviors_merge: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/categories_v021/merge + inherits: behaviors_defaults + type: behavior_class + behaviors_parked_cars: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/categories_v021/parked_cars + inherits: behaviors_defaults + type: behavior_class + behaviors_roundabout: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/categories_v021/roundabout + inherits: behaviors_defaults + type: behavior_class + behaviors_stopped_traffic: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/categories_v021/stopped_traffic + inherits: behaviors_defaults + type: behavior_class + behaviors_traffic_light_green: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/categories_v021/traffic_light_green + inherits: behaviors_defaults + type: behavior_class + behaviors_traffic_light_stop: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/categories_v021/traffic_light_stop + inherits: behaviors_defaults + type: behavior_class + behaviors_unprotected_left: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/categories_v021/unprotected_left + inherits: behaviors_defaults + type: behavior_class + behaviors_unprotected_right: + enabled: 'true' + env: + map_dir: /scratch/ev2237/data/nuplan/categories_v021/unprotected_right + inherits: behaviors_defaults + type: behavior_class + dnf_triage: + enabled: 'false' + env: + map_dir: pufferlib/resources/drive/binaries/carla/opendrive__Town10HD.bin + max_agents_per_env: 1 + min_agents_per_env: 1 + num_maps: 1 + resample_frequency: 500 + scenario_length: 500 + simulation_mode: gigaflow + eval: + num_scenarios: 32 + render_max_steps: 300 + render_num_scenarios: 16 + inherits: validation_defaults + render: 'true' + render_backend: triage_html + type: multi_scenario + validation_defaults: + clean: 'true' + enabled: 'true' + env: + collision_behavior: 1 + eval_mode: 1 + goal_speed: 3.0 + num_agents: 1024 + obs_dropout_boundary: 0.0 + obs_dropout_lane: 0.0 + obs_slots_boundary_n: 80 + obs_slots_lane_n: 80 + offroad_behavior: 1 + reward_ade: 0.0 + reward_collision: 3.0 + reward_comfort: 0.05 + reward_goal: 1.0 + reward_lane_align: 0.025 + reward_lane_center: 0.0038 + reward_offroad: 3.0 + reward_overspeed: 0.05 + reward_randomization: false + reward_reverse: 0.005 + reward_stop_line: 1.0 + reward_timestep: 2.5e-05 + reward_velocity: 0.0025 + goal_source: route + goal_regen_mode: finite + termination_mode: 0 + traffic_light_behavior: 0 + eval: + export_episode_csv: 'true' + num_scenarios: 250 + verify_coverage: 'true' + interval: 250 + mode: inline + validation_gigaflow: + enabled: 'true' + env: + map_dir: pufferlib/resources/drive/binaries/carla + max_agents_per_env: 40 + min_agents_per_env: 40 + num_agents: 1024 + num_maps: 8 + resample_frequency: 500 + scenario_length: 500 + simulation_mode: gigaflow + eval: + render_max_steps: 300 + render_num_scenarios: 8 + inherits: validation_defaults + render: 'true' + render_backend: egl + render_views: + - sim_state + - bev + type: multi_scenario + validation_replay: + enabled: 'true' + env: + control_mode: control_sdc_only + map_dir: /scratch/ev2237/data/nuplan/nuplan_mini_train_bins + max_agents_per_env: 64 + num_maps: 250 + resample_frequency: 200 + scenario_length: 200 + simulation_mode: replay + eval: + render_max_steps: 200 + render_num_scenarios: 5 + inherits: validation_defaults + render: 'true' + render_backend: triage_html + type: multi_scenario + wosac: + clean: 'true' + enabled: 'false' + env: + control_mode: control_wosac + goal_radius: 2.0 + init_mode: create_all_valid + init_step: 10 + eval: + wosac_aggregate_results: 'true' + wosac_num_agents: 256 + wosac_num_rollouts: 32 + wosac_sanity_check: 'false' + interval: 500 + mode: subprocess + render: 'false' + type: wosac +eval_simulation: null +fps: 15 +gif_path: eval.gif +git: + commit_hash: 9d3fcc5c09db50b9674a374b8c3404b6349c4cb5 +load_id: null +load_model_path: null +local_rank: 0 +max_runs: 200 +max_suggestion_cost: 3600 +mine: + num_episodes: 100 + output_dir: '' + render: 'true' + score_threshold: -inf +neptune: false +neptune_name: pufferai +neptune_project: ablations +no_model_upload: {} +num_scenarios: 3 +package: ocean +policy: + actor_hidden_size: 512 + actor_num_layers: 0 + backbone_activation: gelu + backbone_hidden_size: 512 + backbone_layer_norm: false + backbone_num_layers: 2 + boundary_input_size: 128 + context_input_size: 128 + critic_hidden_size: 512 + critic_num_layers: 0 + ego_input_size: 128 + encoder_activation: relu + encoder_layer_norm: true + lane_input_size: 128 + mask_padded_features: false + partner_input_size: 128 + shared_network: true + traffic_control_input_size: 128 + action_type: discrete +policy_name: Drive +render: 0 +render_mode: auto +rnn: + hidden_size: 512 + input_size: 512 +rnn_name: null +run_name: rew-sparse-seed4 +save_frames: 0 +sweep: + downsample: 10 + goal: maximize + method: Protein + metric: score +tag: null +tb: false +train: + adam_beta1: 0.9 + adam_beta2: 0.999 + adam_eps: 1.0e-08 + adv_filter_ewma_beta: 0.25 + adv_filter_threshold_scale: 0.01 + adv_sampling_prio_alpha: 0.8499999999999999 + adv_sampling_prio_beta0: 0.8499999999999999 + amp: true + anneal_lr: true + batch_size: auto + bptt_horizon: 128 + checkpoint_interval: 50 + clip_coef: 0.2 + compile: true + compile_fullgraph: false + compile_mode: default + cpu_offload: false + data_dir: /pufferdrive/training_output/ + device: cuda + ent_coef: 0.01 + gae_lambda: 0.95 + gamma: 0.999 + learning_rate: 0.0005 + max_grad_norm: 0.5 + max_minibatch_size: 98304 + minibatch_size: 98304 + name: pufferai + normalize_rewards: true + obs_only: true + optimizer: adamw + precision: bfloat16 + project: ablations + render: false + render_interval: 1000 + render_map: none + resume_state_path: null + seed: 4 + show_grid: false + show_human_logs: true + show_lasers: false + torch_deterministic: false + total_timesteps: 10000000000 + update_epochs: 2 + use_rnn: false + vf_clip_coef: null + vf_coef: 0.5 + vtrace_c_clip: 1 + vtrace_rho_clip: 1 +vec: + backend: Multiprocessing + batch_size: auto + num_envs: 24 + num_workers: auto + seed: 42 + zero_copy: true +video_path: videos +wandb: true +wandb_group: pr-smoke-runs +wandb_project: pr-smoke-runs diff --git a/tests/unit_tests/test_action_selection.py b/tests/unit_tests/test_action_selection.py new file mode 100644 index 0000000000..fef28938df --- /dev/null +++ b/tests/unit_tests/test_action_selection.py @@ -0,0 +1,196 @@ +"""Unit tests for evaluation-time action selection (sample / mode / mean). + +The feature lets a user pick how eval turns policy logits into actions on a +*continuous* environment, for either a discrete or a continuous policy: + + - sample: draw from the policy distribution (stochastic) + - mode: argmax / Gaussian mean (deterministic "best guess") + - mean: probability-weighted continuous mean (discrete policy only) + +The logic lives in `pufferlib.pytorch.sample_logits`; the mode is chosen by +`Evaluator.action_selection` (config -> validated -> passed to sample_logits). +Both are covered here without an env/C-sim/GPU. +""" + +import sys + +import pytest +import torch + +import pufferlib.pytorch as P + +# 3 discrete classes -> fixed 2D continuous embedding, mirroring Drive's +# action_table[num_classes, 2] (discrete action index -> (long, lat)). +ACTION_TABLE = torch.tensor([[-1.0, -1.0], [0.0, 0.0], [1.0, 1.0]]) + + +class _DiscretePolicy: + """Minimal stand-in for a discrete Drive policy on a continuous env.""" + + is_continuous = False + + def discrete_actions_to_continuous(self, actions): + return ACTION_TABLE[actions.long()] + + def discrete_probs_to_continuous_mean(self, probs): + return probs @ ACTION_TABLE.to(probs.dtype) + + +class _ContinuousPolicy: + is_continuous = True + + +def _discrete_logits(rows): + # Drive's discrete policy returns a tuple of one [batch, num_classes] + # tensor (torch.split output); mirror that so we exercise the real path. + return (torch.tensor(rows),) + + +# -------------------------------------------------------------------------- +# Discrete policy on a continuous env +# -------------------------------------------------------------------------- +def test_mode_selects_argmax_and_is_deterministic(): + # Row 0 argmax = class 2, row 1 argmax = class 0. + logits = _discrete_logits([[0.0, 0.0, 10.0], [10.0, 0.0, 0.0]]) + policy = _DiscretePolicy() + + action1, _, _, cont1 = P.sample_logits( + logits, action_selection=P.ACTION_SELECT_MODE, env_continuous=True, policy=policy + ) + action2, _, _, cont2 = P.sample_logits( + logits, action_selection=P.ACTION_SELECT_MODE, env_continuous=True, policy=policy + ) + + assert action1.reshape(-1).tolist() == [2, 0] + # Continuous action is the exact table row for the argmax class. + expected_cont = torch.stack([ACTION_TABLE[2], ACTION_TABLE[0]]) + assert torch.allclose(cont1.reshape(-1, 2), expected_cont) + # Deterministic: repeated calls are identical. + assert torch.equal(action1, action2) + assert torch.equal(cont1, cont2) + + +def test_mean_is_probability_weighted_and_differs_from_mode(): + # Uniform distribution over the 3 classes. + logits = _discrete_logits([[0.0, 0.0, 0.0]]) + policy = _DiscretePolicy() + + action_mean, _, _, cont_mean = P.sample_logits( + logits, action_selection=P.ACTION_SELECT_MEAN, env_continuous=True, policy=policy + ) + action_mode, _, _, cont_mode = P.sample_logits( + logits, action_selection=P.ACTION_SELECT_MODE, env_continuous=True, policy=policy + ) + + # mean == probability-weighted continuous mean == average of table rows == [0, 0]. + assert torch.allclose(cont_mean.reshape(-1), torch.zeros(2), atol=1e-6) + # mode picks the argmax row ([-1, -1] for a uniform tie) -> distinct from mean. + assert not torch.allclose(cont_mean.reshape(-1), cont_mode.reshape(-1)) + # The nominal discrete action reported for mean is the argmax (for logging). + assert action_mean.reshape(-1).tolist() == action_mode.reshape(-1).tolist() + + +def test_sample_is_stochastic_and_stays_in_support(): + # Flat distribution so multinomial explores every class. + logits = _discrete_logits([[0.0, 0.0, 0.0]]) + policy = _DiscretePolicy() + + torch.manual_seed(0) + seen_classes = set() + for _ in range(50): + action, _, _, cont = P.sample_logits( + logits, action_selection=P.ACTION_SELECT_SAMPLE, env_continuous=True, policy=policy + ) + idx = int(action.reshape(-1).item()) + seen_classes.add(idx) + assert idx in (0, 1, 2) + # Continuous action must be the table row of the *sampled* class. + assert torch.allclose(cont.reshape(-1), ACTION_TABLE[idx]) + # Genuinely sampling, not collapsing to argmax. + assert len(seen_classes) > 1 + + +# -------------------------------------------------------------------------- +# Continuous policy on a continuous env +# -------------------------------------------------------------------------- +def test_continuous_policy_mode_returns_gaussian_mean(): + loc = torch.tensor([[0.5, -0.5]]) + dist = torch.distributions.Normal(loc, torch.ones(1, 2)) + + action, _, _, cont = P.sample_logits( + dist, action_selection=P.ACTION_SELECT_MODE, env_continuous=True, policy=_ContinuousPolicy() + ) + + # A continuous policy already emits a continuous action -> no discrete->cont conversion. + assert cont is None + # Mode of a Gaussian is its mean (loc), deterministically. + assert torch.allclose(action.reshape(1, 2), loc) + + +def test_continuous_policy_sample_differs_from_mode(): + loc = torch.tensor([[0.5, -0.5]]) + dist = torch.distributions.Normal(loc, torch.ones(1, 2)) + + torch.manual_seed(0) + action_sample, _, _, _ = P.sample_logits( + dist, action_selection=P.ACTION_SELECT_SAMPLE, env_continuous=True, policy=_ContinuousPolicy() + ) + + assert not torch.allclose(action_sample.reshape(1, 2), loc) + + +# -------------------------------------------------------------------------- +# Guards: `mean` is only valid for a discrete policy on a continuous env +# -------------------------------------------------------------------------- +def test_mean_rejects_continuous_policy(): + dist = torch.distributions.Normal(torch.zeros(1, 2), torch.ones(1, 2)) + with pytest.raises(ValueError): + P.sample_logits(dist, action_selection=P.ACTION_SELECT_MEAN, env_continuous=True, policy=_ContinuousPolicy()) + + +def test_mean_rejects_discrete_env(): + logits = _discrete_logits([[0.0, 0.0, 0.0]]) + with pytest.raises(ValueError): + P.sample_logits(logits, action_selection=P.ACTION_SELECT_MEAN, env_continuous=False, policy=_DiscretePolicy()) + + +# -------------------------------------------------------------------------- +# Eval config plumbing: eval.action_selection -> validated -> sample_logits +# -------------------------------------------------------------------------- +VALID_ACTION_SELECTIONS = (P.ACTION_SELECT_SAMPLE, P.ACTION_SELECT_MODE, P.ACTION_SELECT_MEAN) + + +def _eval_args(action_selection): + # Minimal args for pufferl.eval; validation runs before any benchmark loading. + return { + "eval": { + "action_selection": action_selection, + "benchmark_config": "pufferlib/config/evaluation/benchmark.yaml", + "benchmarks": None, + "output_name": None, + "render_scenarios": False, + "render_filter": None, + "max_rendered_failures": None, + "failure_replay_csv": None, + "max_sdc_replay_workers": 1, + } + } + + +def test_shipped_config_declares_a_valid_action_selection(monkeypatch): + # Guards the key itself: a config refactor that drops or renames + # eval.action_selection must fail here, not at the first eval run. + import pufferlib.pufferl as pufferl + + # load_config treats everything left in argv as a Hydra override. + monkeypatch.setattr(sys, "argv", ["puffer"]) + args = pufferl.load_config("puffer_drive") + assert args["eval"]["action_selection"] in VALID_ACTION_SELECTIONS + + +def test_eval_rejects_invalid_action_selection(): + import pufferlib + import pufferlib.pufferl as pufferl + + with pytest.raises(pufferlib.APIUsageError): + pufferl.eval(env_name="puffer_drive", args=_eval_args("banana")) diff --git a/weights/mimolette/config.yaml b/weights/mimolette/config.yaml index 3d63068e63..bf8291429c 100644 --- a/weights/mimolette/config.yaml +++ b/weights/mimolette/config.yaml @@ -141,6 +141,7 @@ policy: partner_input_size: 256 shared_network: false traffic_control_input_size: 128 + action_type: discrete policy_name: Drive render: 0 render_mode: matplotlib