Johnny Fernandes
138 lines
4.8 KiB
Python
138 lines
4.8 KiB
Python
"""Observation builder for the shepherd dog policy.
|
|
|
|
Order-invariant 32-D feature vector — the policy generalises across
|
|
flock sizes 1..MAX_SHEEP because individual sheep coordinates never
|
|
appear in the observation by index, only summary statistics, a polar
|
|
histogram, and two "named" sheep (closest-to-pen and rearmost-from-pen).
|
|
|
|
The two named sheep matter for the sequential-driving teacher: it
|
|
targets the closest-to-pen sheep specifically, so the policy needs
|
|
that channel to mimic the teacher.
|
|
|
|
Layout (all components normalised so values stay roughly in [-1, 1]):
|
|
|
|
idx field
|
|
----- ----------------------------------------------------------
|
|
0..3 dog pose: x/15, y/15, cos(heading), sin(heading)
|
|
4..5 active-sheep CoM x/15, y/15
|
|
6..8 flock dispersion: max-radius/15, std_x/15, std_y/15
|
|
9..11 vector dog→CoM: dx/30, dy/30, dist/30
|
|
12..14 vector dog→pen-entry: dx/30, dy/30, dist/30
|
|
15..16 vector furthest-sheep→CoM: dx/15, dy/15
|
|
17..18 min sheep-to-wall, min dog-to-wall (both /15)
|
|
19 active-sheep count / MAX_SHEEP
|
|
20..27 8-bin polar histogram of active sheep around the dog,
|
|
rotation-aware (binned in dog-relative frame), normalised
|
|
so the bins sum to 1.
|
|
28..29 vector dog→closest-to-pen sheep: dx/15, dy/15
|
|
30..31 vector dog→rearmost (furthest-from-pen) sheep: dx/15, dy/15
|
|
"""
|
|
|
|
import math
|
|
import numpy as np
|
|
|
|
from herding.geometry import (
|
|
FIELD_X, FIELD_Y, PEN_ENTRY, MAX_SHEEP,
|
|
)
|
|
|
|
OBS_DIM = 32
|
|
|
|
|
|
def build_obs(dog_xy, dog_heading, sheep_xy_list, sheep_penned_list,
|
|
n_max: int = MAX_SHEEP) -> np.ndarray:
|
|
"""Assemble the dog policy's observation vector.
|
|
|
|
Parameters
|
|
----------
|
|
dog_xy : tuple (x, y) of the dog's GPS position (m)
|
|
dog_heading : dog heading in rad
|
|
sheep_xy_list : iterable of (x, y) for ALL known sheep
|
|
sheep_penned_list : parallel iterable of bool — True if sheep is penned
|
|
n_max : maximum supported flock size used for the count normaliser
|
|
"""
|
|
dog_x, dog_y = dog_xy
|
|
obs = np.zeros(OBS_DIM, dtype=np.float32)
|
|
|
|
obs[0] = dog_x / 15.0
|
|
obs[1] = dog_y / 15.0
|
|
obs[2] = math.cos(dog_heading)
|
|
obs[3] = math.sin(dog_heading)
|
|
|
|
active = [(x, y) for (x, y), p
|
|
in zip(sheep_xy_list, sheep_penned_list) if not p]
|
|
n = len(active)
|
|
|
|
pdx0, pdy0 = PEN_ENTRY[0] - dog_x, PEN_ENTRY[1] - dog_y
|
|
obs[12] = pdx0 / 30.0
|
|
obs[13] = pdy0 / 30.0
|
|
obs[14] = math.hypot(pdx0, pdy0) / 30.0
|
|
|
|
if n == 0:
|
|
# All sheep penned — terminal observation.
|
|
obs[19] = 0.0
|
|
return obs
|
|
|
|
arr = np.asarray(active, dtype=np.float32)
|
|
com_x = float(arr[:, 0].mean())
|
|
com_y = float(arr[:, 1].mean())
|
|
rel = arr - np.array([com_x, com_y], dtype=np.float32)
|
|
dists = np.hypot(rel[:, 0], rel[:, 1])
|
|
radius = float(dists.max())
|
|
std_x = float(arr[:, 0].std())
|
|
std_y = float(arr[:, 1].std())
|
|
|
|
obs[4] = com_x / 15.0
|
|
obs[5] = com_y / 15.0
|
|
obs[6] = radius / 15.0
|
|
obs[7] = std_x / 15.0
|
|
obs[8] = std_y / 15.0
|
|
|
|
cdx, cdy = com_x - dog_x, com_y - dog_y
|
|
obs[9] = cdx / 30.0
|
|
obs[10] = cdy / 30.0
|
|
obs[11] = math.hypot(cdx, cdy) / 30.0
|
|
|
|
far_idx = int(np.argmax(dists))
|
|
obs[15] = float(rel[far_idx, 0]) / 15.0
|
|
obs[16] = float(rel[far_idx, 1]) / 15.0
|
|
|
|
min_sheep_wall = min(
|
|
float(np.min(arr[:, 0] - FIELD_X[0])),
|
|
float(np.min(FIELD_X[1] - arr[:, 0])),
|
|
float(np.min(arr[:, 1] - FIELD_Y[0])),
|
|
float(np.min(FIELD_Y[1] - arr[:, 1])),
|
|
)
|
|
min_dog_wall = min(
|
|
dog_x - FIELD_X[0], FIELD_X[1] - dog_x,
|
|
dog_y - FIELD_Y[0], FIELD_Y[1] - dog_y,
|
|
)
|
|
obs[17] = min_sheep_wall / 15.0
|
|
obs[18] = float(min_dog_wall) / 15.0
|
|
obs[19] = n / n_max
|
|
|
|
# 8-bin polar histogram in the dog's body frame.
|
|
rel_dx = arr[:, 0] - dog_x
|
|
rel_dy = arr[:, 1] - dog_y
|
|
angles = np.arctan2(rel_dy, rel_dx) - dog_heading
|
|
angles = np.arctan2(np.sin(angles), np.cos(angles))
|
|
bins = np.floor((angles + math.pi) / (2 * math.pi) * 8).astype(int)
|
|
bins = np.clip(bins, 0, 7)
|
|
hist = np.bincount(bins, minlength=8).astype(np.float32)
|
|
hist /= max(1, n)
|
|
obs[20:28] = hist
|
|
|
|
# Closest-to-pen sheep (the sequential teacher's target) and rearmost
|
|
# (furthest-from-pen, the natural "next target" once the closest is
|
|
# penned). Both expressed as offset from dog. These two channels make
|
|
# BC tractable — without them the obs doesn't uniquely identify which
|
|
# sheep the teacher is steering toward.
|
|
pen_dists = np.hypot(arr[:, 0] - PEN_ENTRY[0], arr[:, 1] - PEN_ENTRY[1])
|
|
closest_idx = int(np.argmin(pen_dists))
|
|
rearmost_idx = int(np.argmax(pen_dists))
|
|
obs[28] = (float(arr[closest_idx, 0]) - dog_x) / 15.0
|
|
obs[29] = (float(arr[closest_idx, 1]) - dog_y) / 15.0
|
|
obs[30] = (float(arr[rearmost_idx, 0]) - dog_x) / 15.0
|
|
obs[31] = (float(arr[rearmost_idx, 1]) - dog_y) / 15.0
|
|
|
|
return obs
|