TIR_PROJ/herding/world/flocking_sim.py

"""Sheep flocking dynamics — Strömbom 2014 / Reynolds 1987.

Per-sheep behavioural step used by both the Webots sheep controller
and the training environment. Each step a force stack is summed:

    flee       — quadratic ramp away from dog within FLEE_DIST
                 (Strömbom 2014, term ρa)
    cohesion   — drift toward local centre of mass of peers within
                 COHESION_DIST (Strömbom 2014, term c). Weight is
                 higher while fleeing — fear-induced cohesion.
    separation — short-range inverse-distance repulsion from peers
                 (Strömbom 2014 term α; Reynolds 1987)
    wander     — small persistent drift (Strömbom 2014 noise term ε)

Walls, the south-wall gate column, and in-pen containment are
environment-specific additions for the fenced Webots field.

References
----------
- Strömbom et al. (2014). "Solving the shepherding problem: heuristics
  for herding autonomous, interacting agents." J R Soc Interface 11.
- Reynolds (1987). "Flocks, herds and schools: A distributed
  behavioural model." SIGGRAPH '87.
"""

import math
import random

from herding.world.geometry import (
    FIELD_X, FIELD_Y,
    PEN_X, PEN_Y,
    GATE_X,
)

# Speeds are in wheel rad/s (motor units); m/s = speed * SHEEP_WHEEL_RADIUS.
MAX_SPEED = 22.0
FLEE_SPEED = 20.0
WANDER_SPEED = 3.0

WALL_MARGIN = 5.0
WALL_HARD_MARGIN = 1.0
WALL_HARD_GAIN = 50.0

FLEE_DIST = 7.0
SEPARATION_DIST = 2.5
COHESION_DIST = 12.0

PEN_MARGIN = 0.8


def _peers_iter(peers):
    """Accept either a {name: (x, y)} dict or an iterable of (x, y) tuples."""
    if isinstance(peers, dict):
        return list(peers.values())
    return list(peers)


def compute_heading_speed(x, y, penned, dog_xy, peers, wander_angle, rng=None):
    """Return ``(heading, speed, new_wander_angle)`` for one sheep step.

    ``speed`` is in wheel rad/s, bounded by ``[WANDER_SPEED, FLEE_SPEED]``.
    ``heading`` is the world-frame target heading (atan2 convention).
    ``rng`` is an optional ``random.Random`` used for wander jitter; if
    ``None`` uses the module's global ``random``.
    """
    fx, fy = 0.0, 0.0
    peer_list = _peers_iter(peers)
    rnd = rng if rng is not None else random

    if penned:
        # Pen containment: bounce off all four pen walls.
        pm = PEN_MARGIN
        if x < PEN_X[0] + pm:
            fx += ((PEN_X[0] + pm - x) / pm) * 15.0
        if x > PEN_X[1] - pm:
            fx -= ((x - (PEN_X[1] - pm)) / pm) * 15.0
        if y < PEN_Y[0] + pm:
            fy += ((PEN_Y[0] + pm - y) / pm) * 15.0
        if y > PEN_Y[1] - pm:
            fy -= ((y - (PEN_Y[1] - pm)) / pm) * 15.0

        # Mild peer separation so penned sheep don't crowd one corner.
        for px, py in peer_list:
            dx, dy = px - x, py - y
            d = math.hypot(dx, dy)
            if 0.05 < d < SEPARATION_DIST:
                push = (SEPARATION_DIST - d) / d
                fx -= (dx / d) * push * 2.5
                fy -= (dy / d) * push * 2.5

        if rnd.random() < 0.02:
            wander_angle += rnd.uniform(-0.6, 0.6)
        fx += math.cos(wander_angle) * 0.5
        fy += math.sin(wander_angle) * 0.5

    else:
        # Free-roaming sheep in the field.
        fleeing = False
        if dog_xy is not None:
            ddx = dog_xy[0] - x
            ddy = dog_xy[1] - y
            dist = math.hypot(ddx, ddy)
            if 0.01 < dist < FLEE_DIST:
                fleeing = True
                t = 1.0 - dist / FLEE_DIST
                s = t * t * 20.0
                fx -= (ddx / dist) * s
                fy -= (ddy / dist) * s

        # Cohesion: drift toward the local CoM of peers within
        # COHESION_DIST. Stronger while fleeing — fear-induced
        # cohesion keeps the flock together through the gate.
        cx, cy, cn = 0.0, 0.0, 0
        for px, py in peer_list:
            d = math.hypot(px - x, py - y)
            if 0.3 < d < COHESION_DIST:
                cx += px
                cy += py
                cn += 1
        if cn > 0:
            w = 3.0 if fleeing else 1.0
            fx += (cx / cn - x) * w
            fy += (cy / cn - y) * w

        # Separation — inverse-distance push from peers.
        for px, py in peer_list:
            ddx, ddy = px - x, py - y
            d = math.hypot(ddx, ddy)
            if 0.05 < d < SEPARATION_DIST:
                push = (SEPARATION_DIST - d) / d
                fx -= (ddx / d) * push * 2.5
                fy -= (ddy / d) * push * 2.5

        # Wall soft repulsion (south wall absent inside the gate column).
        if x < FIELD_X[0] + WALL_MARGIN:
            fx += ((FIELD_X[0] + WALL_MARGIN - x) / WALL_MARGIN) * 6.0
        if x > FIELD_X[1] - WALL_MARGIN:
            fx -= ((x - (FIELD_X[1] - WALL_MARGIN)) / WALL_MARGIN) * 6.0
        if y > FIELD_Y[1] - WALL_MARGIN:
            fy -= ((y - (FIELD_Y[1] - WALL_MARGIN)) / WALL_MARGIN) * 6.0
        if y < FIELD_Y[0] + WALL_MARGIN and not (GATE_X[0] <= x <= GATE_X[1]):
            fy += ((FIELD_Y[0] + WALL_MARGIN - y) / WALL_MARGIN) * 6.0

        if not fleeing:
            if random.random() < 0.02:
                wander_angle += random.uniform(-0.6, 0.6)
            fx += math.cos(wander_angle) * 0.5
            fy += math.sin(wander_angle) * 0.5

    # Hard escape band — overrides everything else near a wall.
    m, g = WALL_HARD_MARGIN, WALL_HARD_GAIN
    if x - FIELD_X[0] < m:
        fx = max(fx, g * (1.0 - (x - FIELD_X[0]) / m))
    if FIELD_X[1] - x < m:
        fx = min(fx, -g * (1.0 - (FIELD_X[1] - x) / m))
    if FIELD_Y[1] - y < m:
        fy = min(fy, -g * (1.0 - (FIELD_Y[1] - y) / m))
    if (not penned) and (y - FIELD_Y[0] < m) and not (GATE_X[0] <= x <= GATE_X[1]):
        fy = max(fy, g * (1.0 - (y - FIELD_Y[0]) / m))

    heading = math.atan2(fy, fx)
    mag = math.hypot(fx, fy)
    speed = max(WANDER_SPEED, min(FLEE_SPEED, mag * 3.0))
    return heading, speed, wander_angle