Bernhard 01

.gitignore

@@ -221,3 +221,7 @@ docs/_build/
 # Claude config
 .claude/
 CLAUDE.local.md
+
+episode_metrics
+
+/obs

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 = ""
@@ -54,8 +55,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}")
@@ -76,8 +78,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}")
@@ -121,7 +125,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
 
@@ -143,7 +147,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()
@@ -156,8 +162,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
@@ -167,9 +177,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} ---")
@@ -1049,7 +1059,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)

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 = []

notebooks/notebook_utils.py

@@ -112,6 +112,7 @@
     "backbone_layer_norm": False,
     "shared_network": True,
     "mask_padded_features": False,
+    "action_type": "discrete",
 }
 
 

pufferlib/config/ocean/drive.ini

@@ -36,6 +36,8 @@ critic_hidden_size = 512
 critic_num_layers = 0
 ; Dual or shared actor-critic backbone
 shared_network = True
+; Actions type that the policy uses - options: "discrete", "continuous", will be converted to continuous for the simulator.
+action_type = "discrete"
 
 [rnn]
 input_size = 512
@@ -51,7 +53,7 @@ num_agents = 1024
 min_agents_per_env = 1
 max_agents_per_env = 80
 ; 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

pufferlib/config/puffer_drive.yaml

@@ -47,7 +47,7 @@ env:
   min_agents_per_env: 1
   max_agents_per_env: 80
   # 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
@@ -195,6 +195,7 @@ policy:
   critic_num_layers: 0
   # Dual or shared actor-critic backbone
   shared_network: true
+  action_type: discrete
 rnn:
   input_size: 512
   hidden_size: 512
@@ -332,10 +333,15 @@ eval:
       obs_dropout_boundary: 0.0
       obs_slots_lane_n: 80
       obs_slots_boundary_n: 80
+    # 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
     eval:
       num_scenarios: 250
       export_episode_csv: 'true'
       verify_coverage: 'true'
+
   validation_replay:
     inherits: validation_defaults
     type: multi_scenario
@@ -353,6 +359,7 @@ eval:
     eval:
       render_num_scenarios: 5
       render_max_steps: 200
+
   validation_gigaflow:
     inherits: validation_defaults
     type: multi_scenario
@@ -390,6 +397,8 @@ eval:
       num_scenarios: 32
       render_num_scenarios: 16
       render_max_steps: 300
+
+
   # ---------------------------------------------------------------------------
   # Optional: WOSAC realism eval. Off by default.
   # ---------------------------------------------------------------------------
@@ -405,6 +414,8 @@ eval:
       init_mode: create_all_valid
       init_step: 10
       goal_radius: 2.0
+    # Policy action selection during eval: sample | mode | mean.
+    action_selection: mode
     eval:
       wosac_num_rollouts: 32
       wosac_num_agents: 256

pufferlib/ocean/benchmark/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:  # TODO check with trajectory
+                    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:  # TODO check with trajectory
+                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