Checkpoint 2

herding/geometry.py

@@ -0,0 +1,99 @@
+"""World geometry and robot specs.
+
+All coordinates are in meters. (0, 0) is the centre of the field, +x is
+east, +y is north. Z is up but unused here. These constants must match
+``worlds/field.wbt`` and the proto files; if the world changes, change
+this file and only this file.
+
+Pen layout (post-refactor)
+--------------------------
+The pen is *external* to the field, accessed through a 3 m gate cut into
+the south stone wall at y = -15. Sheep entering through the gate end up
+in a fenced rectangle south of the field; the dog stays in the field
+(soft-limited above DOG_SOUTH_LIMIT during training and inference).
+
+    field        +y north
+    +-----------+
+    |           |
+    |           |
+    |  ......   |
+    +---||||----+   y = -15  (south wall, gate at x ∈ [10, 13])
+        ||||
+        |pen|       y ∈ [-22, -15]
+        +---+
+"""
+
+import math
+
+# --- Field (square, stone-walled) ---
+FIELD_X = (-15.0, 15.0)
+FIELD_Y = (-15.0, 15.0)
+
+# Conservative inside bounds — sheep/dog should not graze the wall.
+FIELD_INSIDE_MARGIN = 0.5
+
+# --- Pen (external, south of the field) ---
+PEN_X = (10.0, 13.0)
+PEN_Y = (-22.0, -15.0)
+PEN_CENTER = (0.5 * (PEN_X[0] + PEN_X[1]), 0.5 * (PEN_Y[0] + PEN_Y[1]))
+# The point the dog drives the flock toward: the gate centre on the field side.
+PEN_ENTRY = (0.5 * (PEN_X[0] + PEN_X[1]), -15.0)
+
+# --- Gate (the hole in the south stone wall) ---
+GATE_X = PEN_X
+GATE_Y = -15.0
+
+# --- Robot specs (must match proto files) ---
+# Dog (controllers/shepherd_dog/, protos/ShepherdDog.proto)
+DOG_WHEEL_RADIUS = 0.038         # m
+DOG_WHEEL_BASE = 0.28            # m, axle-to-axle
+DOG_MAX_WHEEL_OMEGA = 70.0       # rad/s
+DOG_MAX_LINEAR = DOG_WHEEL_RADIUS * DOG_MAX_WHEEL_OMEGA  # ~2.66 m/s
+
+# Sheep (controllers/sheep/, protos/Sheep.proto)
+SHEEP_WHEEL_RADIUS = 0.031       # m
+SHEEP_WHEEL_BASE = 0.20          # m
+SHEEP_MAX_WHEEL_OMEGA = 25.0     # rad/s
+SHEEP_MAX_LINEAR = SHEEP_WHEEL_RADIUS * SHEEP_MAX_WHEEL_OMEGA  # ~0.78 m/s
+
+# --- Webots step ---
+WEBOTS_DT = 0.016  # seconds, matches WorldInfo.basicTimeStep = 16 in field.wbt
+
+# --- Dog "virtual south wall" (training keeps dog out of the pen) ---
+# At inference the controller also clips to this so a slightly miscalibrated
+# policy doesn't accidentally drive into the pen and trap the sheep.
+DOG_SOUTH_LIMIT = -14.5
+
+# --- Maximum supported flock size ---
+MAX_SHEEP = 10
+
+
+def in_pen(x: float, y: float) -> bool:
+    """True if (x, y) lies inside the external pen rectangle."""
+    return PEN_X[0] < x < PEN_X[1] and PEN_Y[0] < y < PEN_Y[1]
+
+
+def in_field(x: float, y: float, margin: float = 0.0) -> bool:
+    return (FIELD_X[0] + margin <= x <= FIELD_X[1] - margin
+            and FIELD_Y[0] + margin <= y <= FIELD_Y[1] - margin)
+
+
+def in_gate_corridor(x: float, y: float, margin: float = 0.0) -> bool:
+    """True if (x, y) lies in the column of the gate (between field and pen)."""
+    return (PEN_X[0] - margin <= x <= PEN_X[1] + margin
+            and PEN_Y[0] - margin <= y <= GATE_Y + margin)
+
+
+def is_penned_position(x: float, y: float, latch_margin: float = 0.2) -> bool:
+    """A sheep latches to "penned" once it crosses the gate plane south.
+
+    True iff x is inside the gate column (with a small margin) AND
+    y has dipped below the gate line. Once latched, the sheep is held by
+    in-pen forces and will not exit on its own.
+    """
+    return (PEN_X[0] - latch_margin <= x <= PEN_X[1] + latch_margin
+            and y <= GATE_Y)
+
+
+def distance_to_pen_entry(x: float, y: float) -> float:
+    return math.hypot(x - PEN_ENTRY[0], y - PEN_ENTRY[1])