Files
TIR_PROJ/controllers/sheep/sheep.py
T
Johnny Fernandes 1bb9415414 Checkpoint 2
2026-05-07 22:00:10 +01:00

149 lines
5.1 KiB
Python

"""Sheep flocking controller (Webots).
Each sheep broadcasts its GPS position every 3 steps on channel 1 and
listens for the dog and peer sheep positions. The behavioural step is
delegated to ``herding.flocking_sim.compute_heading_speed`` so the
training environment and Webots run identical sheep dynamics.
Pen behaviour: a sheep latches to ``penned`` the first time it crosses
the south-wall gate plane into the gate corridor. Once latched it turns
pink (via the exposed ``woolColor`` PROTO field) and the force stack
switches to in-pen containment.
"""
import math
import os
import random
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 Supervisor
from herding.diffdrive import heading_speed_to_wheels
from herding.flocking_sim import MAX_SPEED, compute_heading_speed
from herding.geometry import (
SHEEP_MAX_WHEEL_OMEGA,
is_penned_position,
)
# ---------------------------------------------------------------------------
# Device setup
# ---------------------------------------------------------------------------
robot = Supervisor()
timestep = int(robot.getBasicTimeStep())
name = robot.getName()
self_node = robot.getSelf()
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(), SHEEP_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")
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def bearing():
# Compass returns north direction in sensor frame; for this Z-up world
# with north = +Y, atan2(n[0], n[1]) gives the standard math angle
# (0 = east, π/2 = north) matching atan2(fy, fx) used for headings.
n = compass.getValues()
return math.atan2(n[0], n[1])
def drive(heading, speed_motor):
left_w, right_w = heading_speed_to_wheels(
heading, min(speed_motor, MAX_SPEED), bearing(), MOTOR_MAX, k_turn=4.0
)
left_motor.setVelocity(left_w)
right_motor.setVelocity(right_w)
def paint_pink():
# woolColor is declared as a PROTO field with IS binding to the DEF WOOL
# PBRAppearance baseColor; setting it propagates to every USE WOOL shape.
self_node.getField("woolColor").setSFColor([1.0, 0.55, 0.72])
# ---------------------------------------------------------------------------
# State
# ---------------------------------------------------------------------------
wander_angle = random.uniform(-math.pi, math.pi)
step_count = 0
dog_x, dog_y = None, None
peers = {} # name → (x, y), one entry per neighbour, cleared every 30 steps
penned = False
# Stuck detection: differential-drive sheep can pin against a wall and need
# a forced reverse-and-rotate to escape. If displacement < STUCK_DIST for
# STUCK_STEPS consecutive steps, drive toward field centre.
_prev_x, _prev_y = None, None
_stuck_count = 0
STUCK_STEPS = 20
STUCK_DIST = 0.05
# ---------------------------------------------------------------------------
# Main loop
# ---------------------------------------------------------------------------
while robot.step(timestep) != -1:
step_count += 1
pos = gps.getValues()
x, y = pos[0], pos[1]
# Pen entry: one-way latch. Penned sheep get pink wool and switch behaviour.
if not penned and is_penned_position(x, y):
penned = True
paint_pink()
# Refresh peer table — clear before receiving so fresh data is never lost.
if step_count % 30 == 0:
peers.clear()
while receiver.getQueueLength() > 0:
msg = receiver.getString()
receiver.nextPacket()
parts = msg.split(":")
if parts[0] == "dog" and len(parts) >= 3:
dog_x, dog_y = float(parts[1]), float(parts[2])
elif parts[0] == "sheep" and len(parts) >= 4 and parts[1] != name:
peers[parts[1]] = (float(parts[2]), float(parts[3]))
dog_xy = (dog_x, dog_y) if dog_x is not None and dog_y is not None else None
heading, speed, wander_angle = compute_heading_speed(
x=x, y=y, penned=penned, dog_xy=dog_xy, peers=peers,
wander_angle=wander_angle,
)
# Stuck detection — safety net for differential-drive wall pinning.
if _prev_x is not None:
moved = math.hypot(x - _prev_x, y - _prev_y)
_stuck_count = _stuck_count + 1 if moved < STUCK_DIST else 0
if _stuck_count >= STUCK_STEPS:
heading = math.atan2(-y, -x) # always points away from the boundary
speed = MAX_SPEED
_stuck_count = 0
_prev_x, _prev_y = x, y
drive(heading, speed)
if step_count % 3 == 0:
emitter.send(f"sheep:{name}:{x:.4f}:{y:.4f}")