149 lines
5.1 KiB
Python
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}")
|