Allow human to drive with agents through classic and jerk dynamics model (#206)

daphne-cornelisse · Daphne · web-flow · commit b6ed82f80df3 · 2025-12-24T15:29:56.000-05:00
* Fix human control with joint action space &amp; classic model: Was still assuming multi-discrete.

* Enable human control with jerks dynamics model.

* Color actions yellow when controlling.

* Slightly easier control problem?

* Add tiny jerk penalty: Results in smooth behavior.

* Pre-commit

* Minor edits.

* Revert ini changes.

---------

Co-authored-by: Daphne &lt;daphn3cor@gmail.com&gt;
diff --git a/docs/interact-with-agents.md b/docs/interact-with-agents.md
@@ -0,0 +1,46 @@
+## Drive with trained agents
+
+You can take manual control of an agent in the simulator by holding **LEFT SHIFT** and using the keyboard controls. When you're in control, the action values displayed on screen will turn **yellow**.
+
+### Local rendering
+
+To launch an interactive renderer, first build:
+```bash
+bash scripts/build_ocean.sh drive local
+```
+
+then launch:
+```bash
+./drive
+```
+
+This will run `demo()` with an existing model checkpoint.
+
+[TODO: Add demo video/gif here]
+
+### Controls
+
+**General:**
+- **LEFT SHIFT + Arrow Keys/WASD** - Take manual control
+- **TAB** - Switch between agents
+- **SPACE** - First-person camera view
+- **Mouse Drag** - Pan camera
+- **Mouse Wheel** - Zoom
+
+**Classic dynamics model**
+
+- **SHIFT + UP/W** - Increase acceleration
+- **SHIFT + DOWN/S** - Decrease acceleration (brake)
+- **SHIFT + LEFT/A** - Steer left
+- **SHIFT + RIGHT/D** - Steer right
+
+Each key press increments or decrements the action level. For example, tapping W multiple times increases acceleration from neutral (index 3) → 5 → 6 (maximum acceleration). We assume **no friction**, so releasing all keys maintains constant speed and heading.
+
+**Jerk dynamics model**
+
+- **SHIFT + UP/W** - Accelerate (+4.0 m/s³ jerk)
+- **SHIFT + DOWN/S** - Brake (-15.0 m/s³ jerk)
+- **SHIFT + LEFT/A** - Turn left (+4.0 m/s³ lateral jerk)
+- **SHIFT + RIGHT/D** - Turn right (-4.0 m/s³ lateral jerk)
+
+Actions are applied directly when keys are pressed. Pressing W always applies +4.0 m/s³ longitudinal jerk, regardless of how long the key is held.
diff --git a/mkdocs.yml b/mkdocs.yml
@@ -28,6 +28,7 @@ nav:
   - Docs:
       - Getting started: getting-started.md
       - Training agents: train.md
+      - Interact with agents: interact-with-agents.md
       - Simulator: simulator.md
       - Interactive scenario editor: scene-editor.md
       - Visualizer: visualizer.md
diff --git a/pufferlib/ocean/drive/drive.c b/pufferlib/ocean/drive/drive.c
@@ -71,44 +71,65 @@ void demo() {
     int accel_delta = 2;
     int steer_delta = 4;
     while (!WindowShouldClose()) {
-        // Handle camera controls
-        int (*actions)[2] = (int (*)[2])env.actions;
+        int *actions = (int *)env.actions; // Single integer per agent
+
         forward(net, env.observations, env.actions);
+
         if (IsKeyDown(KEY_LEFT_SHIFT)) {
-            actions[env.human_agent_idx][0] = 3;
-            actions[env.human_agent_idx][1] = 6;
-            if (IsKeyDown(KEY_UP) || IsKeyDown(KEY_W)) {
-                actions[env.human_agent_idx][0] += accel_delta;
-                // Cap acceleration to maximum of 6
-                if (actions[env.human_agent_idx][0] > 6) {
-                    actions[env.human_agent_idx][0] = 6;
+            if (env.dynamics_model == CLASSIC) {
+                // Classic dynamics: acceleration and steering
+                int accel_idx = 3; // neutral (0 m/s²)
+                int steer_idx = 6; // neutral (0.0 steering)
+
+                if (IsKeyDown(KEY_UP) || IsKeyDown(KEY_W)) {
+                    accel_idx += accel_delta;
+                    if (accel_idx > 6)
+                        accel_idx = 6;
                 }
-            }
-            if (IsKeyDown(KEY_DOWN) || IsKeyDown(KEY_S)) {
-                actions[env.human_agent_idx][0] -= accel_delta;
-                // Cap acceleration to minimum of 0
-                if (actions[env.human_agent_idx][0] < 0) {
-                    actions[env.human_agent_idx][0] = 0;
+                if (IsKeyDown(KEY_DOWN) || IsKeyDown(KEY_S)) {
+                    accel_idx -= accel_delta;
+                    if (accel_idx < 0)
+                        accel_idx = 0;
                 }
-            }
-            if (IsKeyDown(KEY_LEFT) || IsKeyDown(KEY_A)) {
-                actions[env.human_agent_idx][1] += steer_delta;
-                // Cap steering to minimum of 0
-                if (actions[env.human_agent_idx][1] < 0) {
-                    actions[env.human_agent_idx][1] = 0;
+                if (IsKeyDown(KEY_LEFT) || IsKeyDown(KEY_A)) {
+                    steer_idx += steer_delta; // Increase steering index for left turn
+                    if (steer_idx > 12)
+                        steer_idx = 12;
                 }
-            }
-            if (IsKeyDown(KEY_RIGHT) || IsKeyDown(KEY_D)) {
-                actions[env.human_agent_idx][1] -= steer_delta;
-                // Cap steering to maximum of 12
-                if (actions[env.human_agent_idx][1] > 12) {
-                    actions[env.human_agent_idx][1] = 12;
+                if (IsKeyDown(KEY_RIGHT) || IsKeyDown(KEY_D)) {
+                    steer_idx -= steer_delta; // Decrease steering index for right turn
+                    if (steer_idx < 0)
+                        steer_idx = 0;
                 }
-            }
-            if (IsKeyPressed(KEY_TAB)) {
-                env.human_agent_idx = (env.human_agent_idx + 1) % env.active_agent_count;
+
+                // Encode into single integer: action = accel_idx * 13 + steer_idx
+                actions[env.human_agent_idx] = accel_idx * 13 + steer_idx;
+
+            } else if (env.dynamics_model == JERK) {
+                // Jerk dynamics: longitudinal and lateral jerk
+                // JERK_LONG[4] = {-15.0f, -4.0f, 0.0f, 4.0f}
+                // JERK_LAT[3] = {-4.0f, 0.0f, 4.0f}
+                int jerk_long_idx = 2; // neutral (0.0)
+                int jerk_lat_idx = 1;  // neutral (0.0)
+
+                if (IsKeyDown(KEY_UP) || IsKeyDown(KEY_W)) {
+                    jerk_long_idx = 3; // acceleration (4.0)
+                }
+                if (IsKeyDown(KEY_DOWN) || IsKeyDown(KEY_S)) {
+                    jerk_long_idx = 0; // hard braking (-15.0)
+                }
+                if (IsKeyDown(KEY_LEFT) || IsKeyDown(KEY_A)) {
+                    jerk_lat_idx = 2; // left turn (4.0)
+                }
+                if (IsKeyDown(KEY_RIGHT) || IsKeyDown(KEY_D)) {
+                    jerk_lat_idx = 0; // right turn (-4.0)
+                }
+
+                // Encode into single integer: action = jerk_long_idx * 3 + jerk_lat_idx
+                actions[env.human_agent_idx] = jerk_long_idx * 3 + jerk_lat_idx;
             }
         }
+
         c_step(&env);
         c_render(&env);
     }
diff --git a/pufferlib/ocean/drive/drive.h b/pufferlib/ocean/drive/drive.h
@@ -2020,8 +2020,22 @@ void c_step(Drive *env) {
         env->logs[i].episode_length += 1;
         int agent_idx = env->active_agent_indices[i];
         env->entities[agent_idx].collision_state = 0;
+        float prev_vx = env->entities[agent_idx].vx;
+        float prev_vy = env->entities[agent_idx].vy;
+
         move_dynamics(env, i, agent_idx);
+
+        // Tiny jerk penalty for smoothness
+        if (env->dynamics_model == CLASSIC) {
+            float delta_vx = env->entities[agent_idx].vx - prev_vx;
+            float delta_vy = env->entities[agent_idx].vy - prev_vy;
+            float jerk_penalty = -0.0002f * sqrtf(delta_vx * delta_vx + delta_vy * delta_vy) / env->dt;
+            env->rewards[i] += jerk_penalty;
+            env->logs[i].episode_return += jerk_penalty;
+        }
     }
+
+    // Compute rewards
     for (int i = 0; i < env->active_agent_count; i++) {
         int agent_idx = env->active_agent_indices[i];
         env->entities[agent_idx].collision_state = 0;
@@ -2849,6 +2863,7 @@ void c_render(Drive *env) {
     BeginMode3D(client->camera);
     handle_camera_controls(env->client);
     draw_scene(env, client, 0, 0, 0, 0);
+
     // Draw debug info
     DrawText(TextFormat("Camera Position: (%.2f, %.2f, %.2f)", client->camera.position.x, client->camera.position.y,
                         client->camera.position.z),
@@ -2857,25 +2872,64 @@ void c_render(Drive *env) {
                         client->camera.target.z),
              10, 30, 20, PUFF_WHITE);
     DrawText(TextFormat("Timestep: %d", env->timestep), 10, 50, 20, PUFF_WHITE);
-    // acceleration & steering
+
     int human_idx = env->active_agent_indices[env->human_agent_idx];
     DrawText(TextFormat("Controlling Agent: %d", env->human_agent_idx), 10, 70, 20, PUFF_WHITE);
     DrawText(TextFormat("Agent Index: %d", human_idx), 10, 90, 20, PUFF_WHITE);
-    // Controls help
-    DrawText("Controls: W/S - Accelerate/Brake, A/D - Steer, 1-4 - Switch Agent", 10, client->height - 30, 20,
-             PUFF_WHITE);
-    // acceleration & steering
-    if (env->action_type == 1) { // continuous (float)
+
+    // Display current action values - yellow when controlling, white otherwise
+    Color action_color = IsKeyDown(KEY_LEFT_SHIFT) ? YELLOW : PUFF_WHITE;
+
+    if (env->action_type == 0) { // discrete
+        int *action_array = (int *)env->actions;
+        int action_val = action_array[env->human_agent_idx];
+
+        if (env->dynamics_model == CLASSIC) {
+            int num_steer = 13;
+            int accel_idx = action_val / num_steer;
+            int steer_idx = action_val % num_steer;
+            float accel_value = ACCELERATION_VALUES[accel_idx];
+            float steer_value = STEERING_VALUES[steer_idx];
+
+            DrawText(TextFormat("Acceleration: %.2f m/s^2", accel_value), 10, 110, 20, action_color);
+            DrawText(TextFormat("Steering: %.3f", steer_value), 10, 130, 20, action_color);
+        } else if (env->dynamics_model == JERK) {
+            int num_lat = 3;
+            int jerk_long_idx = action_val / num_lat;
+            int jerk_lat_idx = action_val % num_lat;
+            float jerk_long_value = JERK_LONG[jerk_long_idx];
+            float jerk_lat_value = JERK_LAT[jerk_lat_idx];
+
+            DrawText(TextFormat("Longitudinal Jerk: %.2f m/s^3", jerk_long_value), 10, 110, 20, action_color);
+            DrawText(TextFormat("Lateral Jerk: %.2f m/s^3", jerk_lat_value), 10, 130, 20, action_color);
+        }
+    } else { // continuous
         float (*action_array_f)[2] = (float (*)[2])env->actions;
-        DrawText(TextFormat("Acceleration: %.2f", action_array_f[env->human_agent_idx][0]), 10, 110, 20, PUFF_WHITE);
-        DrawText(TextFormat("Steering: %.2f", action_array_f[env->human_agent_idx][1]), 10, 130, 20, PUFF_WHITE);
-    } else { // discrete (int)
-        int (*action_array)[2] = (int (*)[2])env->actions;
-        DrawText(TextFormat("Acceleration: %d", action_array[env->human_agent_idx][0]), 10, 110, 20, PUFF_WHITE);
-        DrawText(TextFormat("Steering: %d", action_array[env->human_agent_idx][1]), 10, 130, 20, PUFF_WHITE);
-    }
-    DrawText(TextFormat("Grid Rows: %d", env->grid_map->grid_rows), 10, 150, 20, PUFF_WHITE);
-    DrawText(TextFormat("Grid Cols: %d", env->grid_map->grid_cols), 10, 170, 20, PUFF_WHITE);
+        DrawText(TextFormat("Acceleration: %.2f", action_array_f[env->human_agent_idx][0]), 10, 110, 20, action_color);
+        DrawText(TextFormat("Steering: %.2f", action_array_f[env->human_agent_idx][1]), 10, 130, 20, action_color);
+    }
+
+    // Show key press status
+    int status_y = 150;
+    if (IsKeyDown(KEY_LEFT_SHIFT)) {
+        DrawText("[shift pressed]", 10, status_y, 20, YELLOW);
+        status_y += 20;
+    }
+    if (IsKeyDown(KEY_SPACE)) {
+        DrawText("[space pressed]", 10, status_y, 20, YELLOW);
+        status_y += 20;
+    }
+    if (IsKeyDown(KEY_LEFT_CONTROL)) {
+        DrawText("[ctrl pressed]", 10, status_y, 20, YELLOW);
+        status_y += 20;
+    }
+
+    // Controls help
+    DrawText("Controls: SHIFT + W/S - Accelerate/Brake, SHIFT + A/D - Steer, TAB - Switch Agent", 10,
+             client->height - 30, 20, PUFF_WHITE);
+
+    DrawText(TextFormat("Grid Rows: %d", env->grid_map->grid_rows), 10, status_y, 20, PUFF_WHITE);
+    DrawText(TextFormat("Grid Cols: %d", env->grid_map->grid_cols), 10, status_y + 20, 20, PUFF_WHITE);
     EndDrawing();
 }
 
