Johnny Fernandes
276 lines
9.7 KiB
Python
276 lines
9.7 KiB
Python
"""Shepherd Dog controller (Webots).
|
|
|
|
Mode is selected by ``HERDING_MODE`` (env var, or via the
|
|
``herding_runtime.cfg`` file the launcher writes since Webots strips
|
|
env vars on some setups):
|
|
|
|
rl → load a BC-trained SB3 policy from HERDING_POLICY_DIR
|
|
and use its (vx, vy) action each step.
|
|
strombom → canonical Strömbom collect/drive heuristic.
|
|
sequential → single-target "pin and push" — drives the sheep
|
|
closest to the pen.
|
|
|
|
All modes share the same low-level differential-drive controller
|
|
(``herding.diffdrive.velocity_to_wheels`` with cos(err)-clamped forward
|
|
speed), so switching modes does not retune actuation.
|
|
|
|
A safety supervisor enforces the "dog stays out of the pen" invariant:
|
|
if the action would push the dog past ``DOG_SOUTH_LIMIT`` it is
|
|
overridden with a north-driving correction. RL fallback: if the policy
|
|
zip can't be loaded (SB3 missing, file missing), the controller drops
|
|
to strombom mode automatically.
|
|
"""
|
|
|
|
import math
|
|
import os
|
|
import sys
|
|
|
|
# --- Make the shared herding/ package importable from this controller dir ---
|
|
_HERE = os.path.dirname(os.path.abspath(__file__))
|
|
_PROJECT_ROOT = os.path.normpath(os.path.join(_HERE, "..", ".."))
|
|
if _PROJECT_ROOT not in sys.path:
|
|
sys.path.insert(0, _PROJECT_ROOT)
|
|
|
|
from controller import Robot
|
|
|
|
from herding.diffdrive import velocity_to_wheels
|
|
from herding.geometry import (
|
|
DOG_MAX_LINEAR, DOG_MAX_WHEEL_OMEGA,
|
|
DOG_SOUTH_LIMIT, DOG_WHEEL_RADIUS,
|
|
PEN_ENTRY,
|
|
)
|
|
from herding.obs import build_obs
|
|
from herding.sequential import compute_action_debug as sequential_action_debug
|
|
from herding.strombom import compute_action_debug as strombom_action_debug
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Mode selection
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _load_runtime_config():
|
|
"""Read mode + policy_dir overrides from a runtime config file.
|
|
|
|
Webots strips HERDING_* env vars in some configurations, so the
|
|
launcher writes a tiny ``herding_runtime.cfg`` (key=value lines)
|
|
in the project root and the controller reads it here. Env vars
|
|
win if both are present; the file is the fallback.
|
|
"""
|
|
cfg_path = os.path.join(_PROJECT_ROOT, "herding_runtime.cfg")
|
|
if not os.path.exists(cfg_path):
|
|
return {}
|
|
out = {}
|
|
try:
|
|
with open(cfg_path) as f:
|
|
for line in f:
|
|
line = line.strip()
|
|
if not line or line.startswith("#") or "=" not in line:
|
|
continue
|
|
k, _, v = line.partition("=")
|
|
out[k.strip().upper()] = v.strip()
|
|
except OSError:
|
|
return {}
|
|
return out
|
|
|
|
|
|
_runtime_cfg = _load_runtime_config()
|
|
MODE = (os.environ.get("HERDING_MODE")
|
|
or _runtime_cfg.get("HERDING_MODE")
|
|
or "rl").lower()
|
|
|
|
|
|
def _resolve_policy_dir() -> str:
|
|
"""Where to look for the trained policy.
|
|
|
|
Priority:
|
|
1. HERDING_POLICY_DIR env var or runtime-cfg entry, if it points
|
|
to a real directory.
|
|
2. ``training/runs/bc_flock`` — flock-style BC (current default;
|
|
requires the tight-cohesion sheep regime).
|
|
3. ``training/runs/bc_solo`` — single-target BC (1-by-1 style;
|
|
only works if ``herding/flocking_sim.py`` is reverted to the
|
|
loose-cohesion regime).
|
|
"""
|
|
env_dir = (os.environ.get("HERDING_POLICY_DIR")
|
|
or _runtime_cfg.get("HERDING_POLICY_DIR"))
|
|
if env_dir and os.path.isdir(env_dir):
|
|
return env_dir
|
|
candidates = [
|
|
os.path.join(_PROJECT_ROOT, "training", "runs", "bc_flock"),
|
|
os.path.join(_PROJECT_ROOT, "training", "runs", "bc_solo"),
|
|
]
|
|
for c in candidates:
|
|
if os.path.isdir(c):
|
|
return c
|
|
# Last resort — return env var anyway so error message is informative.
|
|
return env_dir or candidates[0]
|
|
|
|
|
|
_VALID_MODES = ("rl", "strombom", "sequential")
|
|
if MODE not in _VALID_MODES:
|
|
print(f"[dog] unknown HERDING_MODE={MODE!r}; defaulting to strombom.")
|
|
MODE = "strombom"
|
|
|
|
POLICY_DIR = _resolve_policy_dir()
|
|
policy_handle = None
|
|
if MODE == "rl":
|
|
print(f"[dog] resolved POLICY_DIR={POLICY_DIR} exists={os.path.isdir(POLICY_DIR)}")
|
|
try:
|
|
from policy_loader import load as _load_policy
|
|
policy_handle = _load_policy(POLICY_DIR)
|
|
print(f"[dog] RL policy loaded from {POLICY_DIR}")
|
|
except Exception as exc:
|
|
print(f"[dog] RL policy load failed ({exc!r}); falling back to strombom.")
|
|
MODE = "strombom"
|
|
print(f"[dog] running in mode={MODE}")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Action smoothing + safety supervisor
|
|
# ---------------------------------------------------------------------------
|
|
|
|
ACTION_SMOOTH = 0.35
|
|
prev_action = (0.0, 0.0)
|
|
|
|
|
|
def safety_clamp(vx: float, vy: float, dog_x: float, dog_y: float) -> tuple:
|
|
"""If the dog is near the south barrier and the action would push it
|
|
further south, override with a northward action. Hard invariant: the
|
|
dog never enters the pen."""
|
|
if dog_y < DOG_SOUTH_LIMIT and vy < 0.0:
|
|
return (0.0, 1.0)
|
|
if dog_y < DOG_SOUTH_LIMIT + 0.5 and vy < -0.2:
|
|
return (vx * 0.5, max(0.0, vy + 0.5))
|
|
return (vx, vy)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Driving
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def drive(vx: float, vy: float, left_motor, right_motor, compass, motor_max: float):
|
|
if math.hypot(vx, vy) < 1e-3:
|
|
left_motor.setVelocity(0.0)
|
|
right_motor.setVelocity(0.0)
|
|
return
|
|
n = compass.getValues()
|
|
h = math.atan2(n[0], n[1])
|
|
left, right = velocity_to_wheels(
|
|
vx, vy, h,
|
|
max_linear=DOG_MAX_LINEAR,
|
|
wheel_radius=DOG_WHEEL_RADIUS,
|
|
max_wheel_omega=motor_max,
|
|
k_turn=4.0,
|
|
)
|
|
left_motor.setVelocity(left)
|
|
right_motor.setVelocity(right)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Webots devices
|
|
# ---------------------------------------------------------------------------
|
|
|
|
robot = Robot()
|
|
timestep = int(robot.getBasicTimeStep())
|
|
|
|
left_motor = robot.getDevice("left wheel motor")
|
|
right_motor = robot.getDevice("right wheel motor")
|
|
left_motor.setPosition(float("inf"))
|
|
right_motor.setPosition(float("inf"))
|
|
left_motor.setVelocity(0.0)
|
|
right_motor.setVelocity(0.0)
|
|
MOTOR_MAX = min(left_motor.getMaxVelocity(), DOG_MAX_WHEEL_OMEGA)
|
|
|
|
gps = robot.getDevice("gps"); gps.enable(timestep)
|
|
compass = robot.getDevice("compass"); compass.enable(timestep)
|
|
receiver = robot.getDevice("receiver"); receiver.enable(timestep)
|
|
emitter = robot.getDevice("emitter")
|
|
|
|
# Cosmetic ear motors — ignored by control logic but keep them animated.
|
|
left_ear = robot.getDevice("left ear motor")
|
|
right_ear = robot.getDevice("right ear motor")
|
|
left_ear.setPosition(float("inf"))
|
|
right_ear.setPosition(float("inf"))
|
|
left_ear.setVelocity(0.0)
|
|
right_ear.setVelocity(0.0)
|
|
ear_phase = 0.0
|
|
EAR_AMPLITUDE = 0.35
|
|
EAR_RATE = 8.0
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Main loop
|
|
# ---------------------------------------------------------------------------
|
|
|
|
# {name: (x, y)} — kept across all sheep ever heard from. Sheep that drift
|
|
# into the pen are tracked by ``penned`` so observations and Strömbom
|
|
# agree on which ones still need herding.
|
|
sheep_positions: dict = {}
|
|
penned_set: set = set()
|
|
step_count = 0
|
|
|
|
from herding.geometry import is_penned_position
|
|
|
|
while robot.step(timestep) != -1:
|
|
step_count += 1
|
|
|
|
while receiver.getQueueLength() > 0:
|
|
msg = receiver.getString()
|
|
receiver.nextPacket()
|
|
parts = msg.split(":")
|
|
if len(parts) == 4 and parts[0] == "sheep":
|
|
try:
|
|
x, y = float(parts[2]), float(parts[3])
|
|
except ValueError:
|
|
continue
|
|
sheep_positions[parts[1]] = (x, y)
|
|
if parts[1] not in penned_set and is_penned_position(x, y):
|
|
penned_set.add(parts[1])
|
|
|
|
pos = gps.getValues()
|
|
dog_xy = (pos[0], pos[1])
|
|
n = compass.getValues()
|
|
dog_heading = math.atan2(n[0], n[1])
|
|
|
|
# ---- Action selection ----
|
|
if MODE == "rl" and policy_handle is not None:
|
|
sheep_xy_list = list(sheep_positions.values())
|
|
sheep_names = list(sheep_positions.keys())
|
|
sheep_penned_list = [s in penned_set for s in sheep_names]
|
|
obs = build_obs(dog_xy, dog_heading, sheep_xy_list, sheep_penned_list)
|
|
action = policy_handle.predict(obs)
|
|
vx, vy = float(action[0]), float(action[1])
|
|
elif MODE == "sequential":
|
|
vx, vy, _mode_str, _dbg = sequential_action_debug(
|
|
dog_xy, sheep_positions, PEN_ENTRY,
|
|
)
|
|
else:
|
|
# Strömbom (canonical baseline).
|
|
vx, vy, _mode_str, _dbg = strombom_action_debug(
|
|
dog_xy, sheep_positions, PEN_ENTRY,
|
|
)
|
|
|
|
# EMA smoothing — reduces oscillation from policy or Strömbom flips.
|
|
vx = ACTION_SMOOTH * prev_action[0] + (1.0 - ACTION_SMOOTH) * vx
|
|
vy = ACTION_SMOOTH * prev_action[1] + (1.0 - ACTION_SMOOTH) * vy
|
|
|
|
# Safety: dog must never enter the pen.
|
|
vx, vy = safety_clamp(vx, vy, dog_xy[0], dog_xy[1])
|
|
prev_action = (vx, vy)
|
|
|
|
drive(vx, vy, left_motor, right_motor, compass, MOTOR_MAX)
|
|
emitter.send(f"dog:{dog_xy[0]:.4f}:{dog_xy[1]:.4f}")
|
|
|
|
# Cosmetic ear wiggle — purely visual.
|
|
ear_phase += 0.12
|
|
ear_pos = EAR_AMPLITUDE * math.sin(ear_phase)
|
|
left_ear.setVelocity(EAR_RATE)
|
|
right_ear.setVelocity(EAR_RATE)
|
|
left_ear.setPosition(ear_pos)
|
|
right_ear.setPosition(-ear_pos)
|
|
|
|
if step_count % 200 == 0:
|
|
n_active = sum(1 for s in sheep_positions if s not in penned_set)
|
|
print(f"[dog mode={MODE}] step={step_count} known={len(sheep_positions)} "
|
|
f"penned={len(penned_set)} active={n_active} action=({vx:+.2f}, {vy:+.2f})")
|