TIR_PROJ/controllers/shepherd_dog/shepherd_dog.py

"""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):

    strombom    → canonical Strömbom (2014) collect/drive heuristic
                  wrapped in ActiveScanTeacher (opening rotation +
                  walk-to-centre when the tracker briefly empties).
    sequential  → single-target "pin-and-push", same wrapper.
    bc          → behaviour-cloned MLP, trained on Strömbom demos.
                  Default policy: training/runs/bc/policy.zip.
    rl          → KL-regularised PPO fine-tune of bc. Same obs/action
                  space as bc; refines time-to-pen via reward while
                  staying anchored to bc.
                  Default policy: training/runs/rl/policy.zip.

Sheep perception
----------------
The dog perceives sheep through its **front-mounted 140° LiDAR**
(``protos/ShepherdDog.proto``: 180 rays, 12 m max range). Each step:

    1. Reads ``lidar.getRangeImage()``.
    2. Runs ``herding.perception.lidar_perception.detections_from_scan``
       to cluster returns into world-frame ``(x, y)`` sheep estimates.
    3. Folds those into a ``SheepTracker`` which maintains last-seen
       positions for sheep currently out of FOV and latches "penned"
       once a track crosses the gate plane south.

Sheep ``emitter`` messages are read **for diagnostic logging only**
(GT_penned counter + auto-finish sentinel); they are never used to
drive the policy. Perception for control comes entirely from LiDAR.

Auto-finish
-----------
When the dog observes (via GT, read off the receiver) that all sheep
are penned, it writes ``training/.run_done`` and the launcher
(``tools/run_webots.sh``) detects it and closes Webots. This keeps
batch evaluation runs bounded.
"""

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)

# --- Read runtime cfg early so env vars are set before geometry import ---
def _load_runtime_config():
    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()
# Seed env vars from runtime cfg so downstream modules (geometry.py) see them.
for _rk, _rv in _runtime_cfg.items():
    if _rk.startswith("HERDING_") and _rk not in os.environ:
        os.environ[_rk] = _rv

import numpy as np

from controller import Robot

from herding.control.active_scan import ActiveScanTeacher
from herding.control.modulation import modulate_speed_near_sheep
from herding.control.sequential import compute_action as sequential_action
from herding.control.strombom import compute_action as strombom_action
from herding.control.universal import compute_action as universal_action
from herding.perception.obs import build_obs
from herding.perception.lidar_perception import detections_from_scan
from herding.perception.sheep_tracker import SheepTracker
from herding.world.diffdrive import velocity_to_mecanum_wheels, velocity_to_wheels
from herding.world.geometry import (
    DOG_SOUTH_LIMIT,
    PEN_ENTRY, is_penned_position,
)
from herding.config import HERDING_WEBOTS, RobotConfig

# Robot physical constants come from RobotConfig so they stay in sync with
# the training environment.  The Webots preset uses action_smooth=0.55.
_ROBOT_CFG: RobotConfig = HERDING_WEBOTS.robot
DOG_WHEEL_RADIUS  = _ROBOT_CFG.wheel_radius
DOG_WHEEL_BASE    = _ROBOT_CFG.wheel_base
DOG_WHEEL_BASE_X  = _ROBOT_CFG.wheel_base_x
DOG_WHEEL_BASE_Y  = _ROBOT_CFG.wheel_base_y
DOG_MAX_WHEEL_OMEGA = _ROBOT_CFG.max_wheel_omega
DOG_MAX_LINEAR    = _ROBOT_CFG.max_linear


# ---------------------------------------------------------------------------
# Mode + policy resolution (cfg already loaded above)
# ---------------------------------------------------------------------------

MODE = (os.environ.get("HERDING_MODE")
        or _runtime_cfg.get("HERDING_MODE")
        or "bc").lower()

# Diagnostic: bypass LiDAR tracker and use GT emitter positions directly.
# Set HERDING_USE_GT=1 to isolate perception sim-to-real gap from policy quality.
USE_GT_PERCEPTION = bool(int(
    os.environ.get("HERDING_USE_GT")
    or _runtime_cfg.get("HERDING_USE_GT", "0")
))


def _resolve_policy_dir(mode: str) -> str:
    """Where to look for the trained policy for the given mode.

    Priority:
      1. HERDING_POLICY_DIR env var or runtime-cfg entry, if it points
         to a real directory.
      2. Drive-mode-specific default:
            bc → training/runs/bc_differential (or bc_mecanum)
            rl → training/runs/rl_differential (or rl_mecanum)
      3. Legacy path (no drive suffix):
            bc → training/runs/bc
            rl → training/runs/rl
    """
    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
    drive = DRIVE_MODE
    mode_default = {
        "bc": os.path.join(_PROJECT_ROOT, "training", "runs",
                           f"bc_{drive}"),
        "rl": os.path.join(_PROJECT_ROOT, "training", "runs",
                           f"rl_{drive}"),
    }
    primary = mode_default.get(mode, mode_default["bc"])
    if os.path.isdir(primary):
        return primary
    # Fallback: legacy paths without drive suffix.
    legacy = {
        "bc": os.path.join(_PROJECT_ROOT, "training", "runs", "bc"),
        "rl": os.path.join(_PROJECT_ROOT, "training", "runs", "rl"),
    }
    fallback = legacy.get(mode, legacy["bc"])
    if os.path.isdir(fallback):
        return fallback
    return env_dir or primary


_VALID_MODES = ("bc", "rl", "strombom", "sequential", "universal", "calibrate")
if MODE not in _VALID_MODES:
    print(f"[dog] unknown HERDING_MODE={MODE!r}; defaulting to strombom.")
    MODE = "strombom"

POLICY_DIR = _resolve_policy_dir(MODE)
policy_handle = None
if MODE in ("bc", "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] policy loaded from {POLICY_DIR}")
    except Exception as exc:
        print(f"[dog] policy load failed ({exc!r}); falling back to strombom.")
        MODE = "strombom"
print(f"[dog] running in mode={MODE}")

# Drive mode: "differential" (2-wheel) or "mecanum" (4-wheel omnidirectional).
DRIVE_MODE = (os.environ.get("HERDING_DRIVE")
              or _runtime_cfg.get("HERDING_DRIVE")
              or "differential").lower()
if DRIVE_MODE not in ("differential", "mecanum"):
    print(f"[dog] unknown HERDING_DRIVE={DRIVE_MODE!r}; defaulting to differential.")
    DRIVE_MODE = "differential"
print(f"[dog] drive mode={DRIVE_MODE}")


# ---------------------------------------------------------------------------
# Control parameters
# ---------------------------------------------------------------------------

ACTION_SMOOTH = _ROBOT_CFG.action_smooth  # EMA on (vx, vy) — kills frame-to-frame jitter
RUN_DONE_FILE = os.path.join(_PROJECT_ROOT, "training", ".run_done")


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)


def drive_diff(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)


def drive_mecanum(vx: float, vy: float, omega: float,
                  fl_motor, fr_motor, rl_motor, rr_motor,
                  compass, motor_max: float):
    if math.hypot(vx, vy) < 1e-3 and abs(omega) < 1e-3:
        fl_motor.setVelocity(0.0)
        fr_motor.setVelocity(0.0)
        rl_motor.setVelocity(0.0)
        rr_motor.setVelocity(0.0)
        return
    n = compass.getValues()
    h = math.atan2(n[0], n[1])
    w_fl, w_fr, w_rl, w_rr = velocity_to_mecanum_wheels(
        vx, vy, omega, h,
        max_linear=DOG_MAX_LINEAR,
        wheel_radius=DOG_WHEEL_RADIUS,
        lx=DOG_WHEEL_BASE_X / 2.0, ly=DOG_WHEEL_BASE_Y / 2.0,
        max_wheel_omega=motor_max,
        k_turn=4.0,
        wheel_base=DOG_WHEEL_BASE,
    )
    fl_motor.setVelocity(w_fl)
    fr_motor.setVelocity(w_fr)
    rl_motor.setVelocity(w_rl)
    rr_motor.setVelocity(w_rr)


# ---------------------------------------------------------------------------
# Webots devices
# ---------------------------------------------------------------------------

robot = Robot()
timestep = int(robot.getBasicTimeStep())

if DRIVE_MODE == "mecanum":
    fl_motor = robot.getDevice("front left wheel motor")
    fr_motor = robot.getDevice("front right wheel motor")
    rl_motor = robot.getDevice("rear left wheel motor")
    rr_motor = robot.getDevice("rear right wheel motor")
    for m in (fl_motor, fr_motor, rl_motor, rr_motor):
        m.setPosition(float("inf"))
        m.setVelocity(0.0)
    MOTOR_MAX = min(fl_motor.getMaxVelocity(), DOG_MAX_WHEEL_OMEGA)
else:
    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")
lidar = robot.getDevice("lidar");       lidar.enable(timestep)

tracker = SheepTracker(tracker_cfg=HERDING_WEBOTS.tracker)

# Cosmetic ear motors — animated; not used by control.
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
# ---------------------------------------------------------------------------

# Analytic-teacher wrapper (instantiated lazily so RL/BC modes don't pay
# the import-time cost). Each gets the same ActiveScanTeacher treatment:
# rotate-on-empty, walk-to-centre, near-sheep speed modulation.
analytic_teacher = None
if MODE in ("strombom", "sequential"):
    base_fn = strombom_action if MODE == "strombom" else sequential_action
    analytic_teacher = ActiveScanTeacher(base_fn)
elif MODE == "universal":
    analytic_teacher = ActiveScanTeacher(universal_action)

# GT positions from sheep emitters — used **only** for the auto-finish
# sentinel and the GT_penned diagnostic line. Never fed into control.
_gt_sheep: dict = {}
_run_done = False

prev_action = (0.0, 0.0, 0.0) if DRIVE_MODE == "mecanum" else (0.0, 0.0)
step_count = 0

# ---------------------------------------------------------------------------
# Calibration mode — apply fixed action, measure GPS displacement, compare
# against gym kinematics prediction, write results to calibrate_mecanum.log.
# ---------------------------------------------------------------------------
if MODE == "calibrate":
    import json as _json
    _calib_vx  = float(os.environ.get("CALIB_VX", "0.5"))
    _calib_vy  = float(os.environ.get("CALIB_VY", "0.0"))
    _calib_om  = float(os.environ.get("CALIB_OM", "0.0"))
    _calib_n   = int(os.environ.get("CALIB_N_STEPS", "150"))
    _log_path  = os.path.join(_PROJECT_ROOT, "calibrate_mecanum.log")
    # Settle for 5 steps so GPS stabilises.
    for _ in range(5):
        robot.step(timestep)
    _pos0 = gps.getValues(); _x0, _y0 = _pos0[0], _pos0[1]
    _n_calib = compass.getValues(); _h0 = math.atan2(_n_calib[0], _n_calib[1])
    # Gym-predicted displacement using shared kinematics.
    from herding.world.diffdrive import velocity_to_mecanum_wheels, mecanum_kinematics_step
    from herding.world.geometry import WEBOTS_DT as _DT
    _xg, _yg, _hg = _x0, _y0, _h0
    for _ in range(_calib_n):
        _wfl, _wfr, _wrl, _wrr = velocity_to_mecanum_wheels(
            _calib_vx, _calib_vy, _calib_om, _hg,
            max_linear=DOG_MAX_LINEAR, wheel_radius=DOG_WHEEL_RADIUS,
            lx=DOG_WHEEL_BASE_X / 2, ly=DOG_WHEEL_BASE_Y / 2,
            max_wheel_omega=DOG_MAX_WHEEL_OMEGA, k_turn=4.0,
            wheel_base=DOG_WHEEL_BASE,
        )
        _xg, _yg, _hg = mecanum_kinematics_step(
            _xg, _yg, _hg, _wfl, _wfr, _wrl, _wrr,
            DOG_WHEEL_RADIUS, DOG_WHEEL_BASE_X / 2, DOG_WHEEL_BASE_Y / 2, _DT,
        )
    # Run actual Webots steps.
    for _ in range(_calib_n):
        _nv = compass.getValues(); _h = math.atan2(_nv[0], _nv[1])
        _wfl, _wfr, _wrl, _wrr = velocity_to_mecanum_wheels(
            _calib_vx, _calib_vy, _calib_om, _h,
            max_linear=DOG_MAX_LINEAR, wheel_radius=DOG_WHEEL_RADIUS,
            lx=DOG_WHEEL_BASE_X / 2, ly=DOG_WHEEL_BASE_Y / 2,
            max_wheel_omega=DOG_MAX_WHEEL_OMEGA, k_turn=4.0,
            wheel_base=DOG_WHEEL_BASE,
        )
        if DRIVE_MODE == "mecanum":
            drive_mecanum(_calib_vx, _calib_vy, _calib_om,
                          fl_motor, fr_motor, rl_motor, rr_motor,
                          compass, MOTOR_MAX)
        robot.step(timestep)
    _pos1 = gps.getValues(); _x1, _y1 = _pos1[0], _pos1[1]
    _T = _calib_n * _DT
    _vx_w = (_x1 - _x0) / _T;  _vy_w = (_y1 - _y0) / _T
    _vx_g = (_xg - _x0) / _T;  _vy_g = (_yg - _y0) / _T
    def _pct(a, p): return 0.0 if abs(p) < 1e-4 else 100.0 * abs(a - p) / abs(p)
    _result = (
        f"cmd=(vx={_calib_vx:.2f}, vy={_calib_vy:.2f}, om={_calib_om:.2f})  "
        f"steps={_calib_n}\n"
        f"  gym    displacement: dx={_xg-_x0:.4f}  dy={_yg-_y0:.4f}  "
        f"(vx={_vx_g:.3f}  vy={_vy_g:.3f} m/s)\n"
        f"  webots displacement: dx={_x1-_x0:.4f}  dy={_y1-_y0:.4f}  "
        f"(vx={_vx_w:.3f}  vy={_vy_w:.3f} m/s)\n"
        f"  vx error: {_pct(_vx_w, _vx_g):.1f}%    "
        f"vy error: {_pct(_vy_w, _vy_g):.1f}%\n"
    )
    print(_result)
    with open(_log_path, "a") as _f:
        _f.write(_result + "\n")
    # Write the run-done sentinel so run_webots.sh closes Webots cleanly.
    with open(RUN_DONE_FILE, "w") as _f:
        _f.write("calibrate\n")
    import sys as _sys; _sys.exit(0)

while robot.step(timestep) != -1:
    step_count += 1

    # Drain sheep emitter messages → GT (diagnostic only).
    while receiver.getQueueLength() > 0:
        msg = receiver.getString()
        receiver.nextPacket()
        parts = msg.split(":")
        if len(parts) == 4 and parts[0] == "sheep":
            try:
                _gt_sheep[parts[1]] = (float(parts[2]), float(parts[3]))
            except ValueError:
                pass

    pos = gps.getValues()
    dog_xy = (pos[0], pos[1])
    n = compass.getValues()
    dog_heading = math.atan2(n[0], n[1])

    # ---- LiDAR perception → tracker → active sheep positions ----
    ranges = np.asarray(lidar.getRangeImage(), dtype=np.float32)
    detections = detections_from_scan(
        ranges, dog_xy[0], dog_xy[1], dog_heading,
        detection_cfg=HERDING_WEBOTS.detection,
        lidar_cfg=HERDING_WEBOTS.lidar,
    )
    if USE_GT_PERCEPTION and _gt_sheep:
        # Bypass tracker: feed GT emitter positions directly to policy/teacher.
        sheep_positions = {k: v for k, v in _gt_sheep.items()
                          if not is_penned_position(v[0], v[1])}
        tracker.update(detections)  # still advance tracker for diagnostics
    else:
        sheep_positions = tracker.update(detections)

    sheep_xy_list = list(sheep_positions.values())
    sheep_penned_list = [False] * len(sheep_xy_list)
    single_obs = build_obs(dog_xy, dog_heading, sheep_xy_list, sheep_penned_list)

    # ---- Action selection ----
    omega = 0.0
    if MODE in ("bc", "rl") and policy_handle is not None:
        if not sheep_positions:
            # BC/RL never saw "empty obs during operation" in training (empty
            # obs only happened at episode end), so the policy outputs ~zero
            # and the dog gets stuck. Fall back to a fixed scan rotation
            # until tracker recovers some sheep.
            vx, vy = 0.0, 0.6
            omega = 0.5 if DRIVE_MODE == "mecanum" else 0.0
        else:
            action = policy_handle.predict(single_obs)
            vx, vy = float(action[0]), float(action[1])
            if DRIVE_MODE == "mecanum" and len(action) >= 3:
                omega = float(action[2])
    else:
        result = analytic_teacher(
            dog_xy, dog_heading, sheep_positions, PEN_ENTRY,
            DRIVE_MODE,
        )
        if len(result) == 4:
            vx, vy, omega, _mode_str = result
        else:
            vx, vy, _mode_str = result

    # Near-sheep speed modulation (shared by every mode).
    vx, vy = modulate_speed_near_sheep(vx, vy, dog_xy, sheep_positions)

    # EMA smoothing — kills frame-to-frame action jitter.
    if DRIVE_MODE == "mecanum":
        vx = ACTION_SMOOTH * prev_action[0] + (1.0 - ACTION_SMOOTH) * vx
        vy = ACTION_SMOOTH * prev_action[1] + (1.0 - ACTION_SMOOTH) * vy
        omega = ACTION_SMOOTH * prev_action[2] + (1.0 - ACTION_SMOOTH) * omega
    else:
        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, omega) if DRIVE_MODE == "mecanum" else (vx, vy)

    if DRIVE_MODE == "mecanum":
        drive_mecanum(vx, vy, omega, fl_motor, fr_motor, rl_motor, rr_motor,
                      compass, MOTOR_MAX)
    else:
        drive_diff(vx, vy, left_motor, right_motor, compass, MOTOR_MAX)
    emitter.send(f"dog:{dog_xy[0]:.4f}:{dog_xy[1]:.4f}")

    # Cosmetic ear wiggle.
    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)

    # Auto-finish: when all GT sheep are penned, write the sentinel.
    # The launcher polls for it and closes Webots so batch evals don't
    # hang after the task is done. Bounded by `_gt_sheep` so we don't
    # fire during the first few steps while the receiver fills.
    if _gt_sheep and not _run_done:
        gt_active = sum(1 for x, y in _gt_sheep.values()
                        if not is_penned_position(x, y))
        if gt_active == 0:
            os.makedirs(os.path.dirname(RUN_DONE_FILE), exist_ok=True)
            open(RUN_DONE_FILE, "w").close()
            _run_done = True
            print(f"[dog] all {len(_gt_sheep)} sheep penned at step "
                  f"{step_count} — wrote sentinel, launcher will close Webots")

    if step_count % 200 == 0:
        gt_penned = sum(1 for x, y in _gt_sheep.values()
                        if is_penned_position(x, y))
        gt_total = len(_gt_sheep)
        print(f"[dog mode={MODE} drive={DRIVE_MODE}] step={step_count} "
              f"GT_penned={gt_penned}/{gt_total} "
              f"tracks_active={tracker.n_active()} "
              f"tracks_penned={tracker.n_penned()} "
              f"detections={len(detections)} "
              f"action=({vx:+.2f}, {vy:+.2f}, {omega:+.2f})"
              if DRIVE_MODE == "mecanum" else
              f"[dog mode={MODE} drive={DRIVE_MODE}] step={step_count} "
              f"GT_penned={gt_penned}/{gt_total} "
              f"tracks_active={tracker.n_active()} "
              f"tracks_penned={tracker.n_penned()} "
              f"detections={len(detections)} action=({vx:+.2f}, {vy:+.2f})")