Sheep training flock of 10 fix?

controllers/shepherd_dog_rl/shepherd_dog_rl.py

@@ -80,14 +80,15 @@ def build_obs(dog_pos: np.ndarray,
     n_active = len(active_pos)
 
     if n_active > 0:
-        com       = active_pos.mean(axis=0)
+        com        = active_pos.mean(axis=0)
         d_from_com = np.linalg.norm(active_pos - com, axis=1)
-        radius    = float(d_from_com.max())
-        mean_disp = float(d_from_com.mean())
-        far       = active_pos[int(np.argmax(d_from_com))]
+        sorted_idx = np.argsort(d_from_com)[::-1]
+        radius     = float(d_from_com[sorted_idx[0]])
+        far        = active_pos[sorted_idx[0]]
+        second_far_dist = float(d_from_com[sorted_idx[1]]) if len(sorted_idx) > 1 else 0.0
     else:
         com = PEN_CENTER.copy()
-        radius = mean_disp = 0.0
+        radius = second_far_dist = 0.0
         far = PEN_CENTER.copy()
 
     frac_active = n_active / max(n_sheep, 1)
@@ -98,8 +99,8 @@ def build_obs(dog_pos: np.ndarray,
         (far[0] - dog_pos[0]) / D,  (far[1] - dog_pos[1]) / D,
         (PEN_CENTER[0] - com[0]) / D,  (PEN_CENTER[1] - com[1]) / D,
         (PEN_CENTER[0] - far[0]) / D,  (PEN_CENTER[1] - far[1]) / D,
-        radius    / D,
-        mean_disp / D,
+        radius          / D,
+        second_far_dist / D,
         frac_active,
     ], dtype=np.float32)