Mimic webots physics

controllers/shepherd_dog_rl/shepherd_dog_rl.py

@@ -79,11 +79,13 @@ def in_pen(x: float, y: float) -> bool:
 
 def build_obs(dog_pos: np.ndarray,
               sheep_dict: dict,
-              n_sheep: int) -> np.ndarray:
+              n_sheep: int,
+              dog_heading: float = 0.0) -> np.ndarray:
     """
-    Build the 13-dim flock observation — identical to HerdingEnv._obs().
+    Build the 18-dim flock observation — identical to HerdingEnv._obs().
 
     sheep_dict: {name: (x, y)} for ALL known sheep (penned or not).
+    dog_heading: dog's current world-frame heading in radians.
     """
     D = 2 * FIELD
 
@@ -119,6 +121,7 @@ def build_obs(dog_pos: np.ndarray,
         (PEN_CENTER[0] - far1[0]) / D, (PEN_CENTER[1] - far1[1]) / D,
         radius / D,
         frac_active,
+        math.cos(dog_heading), math.sin(dog_heading),
     ], dtype=np.float32)
 
 
@@ -152,7 +155,7 @@ ear_phase = 0.0
 try:
     n_sheep = int(sys.argv[1])
 except (IndexError, ValueError):
-    n_sheep = 5
+    n_sheep = 3
 
 # ── Load model ───────────────────────────────────────────────────────────────
 print(f"[RL dog] Loading model from {MODEL_PATH}")
@@ -230,8 +233,9 @@ while robot.step(timestep) != -1:
     gps_vals = gps.getValues()
     dog_pos  = np.array([gps_vals[0], gps_vals[1]], dtype=np.float32)
 
-    # 3. Build and normalise observation
-    raw_obs  = build_obs(dog_pos, sheep_positions, n_sheep)
+    # 3. Build and normalise observation (heading from compass)
+    raw_obs  = build_obs(dog_pos, sheep_positions, n_sheep,
+                         dog_heading=bearing())
     obs_norm = vecnorm.normalize_obs(raw_obs[np.newaxis])  # (1, 13)
 
     # 4. Policy inference + smoothing