Files
TIR_PROJ/training/bc/collect.py
T
Johnny Fernandes dd5ac669e5 Webots sim-to-real fixes, DAgger pipeline, 360° proto variant
Today's session worked across the full Webots delivery stack — found and
fixed a cluster of bugs blocking the BC/RL transfer, then explored
training-side mitigations for the residual perception gap.

Bug fixes:
- Makefile FP_RATE default 2.0 → 0.0: BC demos used fp_rate=0 but RL
  fine-tune defaulted to fp_rate=2, poisoning the BC obs distribution
  and stalling PPO at 0% success across 1.46M+ steps.
- controllers/{shepherd_dog,sheep}/runtime.ini: Webots was launching
  controllers under system python3 (no numpy) and they were crashing
  silently. Pinned to the conda tir env.
- herding/config.py HERDING_WEBOTS preset: pen_latch_depth 0.5 → 2.0,
  max_new_tracks_per_step 3 → 1, static_reject 0.8 → 1.2. Stops phantom
  FPs near the gate from latching as permanently-penned tracks.
- herding/perception/sheep_tracker.py: penned tracks now decay at
  forget_steps × 8 instead of living forever. Adds get_positions
  min_freshness filter for deploy-time use.

Training/eval matches deployment:
- training/bc/collect.py: --dagger-policy flag for DAgger rollouts
  (policy drives, teacher labels) + --use-webots-preset for matched
  140° tracker + DR config.
- controllers/shepherd_dog/shepherd_dog.py: scan-fallback (0, 0.6) when
  BC/RL sees empty sheep_positions — recovers from FOV gaps.

Tooling:
- tools/dagger_round.sh: one-shot DAgger round (collect + concat + bc).
- tools/webots_sweep_gt.sh: full sweep with HERDING_USE_GT=1 for the
  perception-gap diagnosis matrix.
- protos/ShepherdDog360.proto: 360° FOV variant for the FOV-ablation
  comparison. Canonical proto stays at 140° per project spec.

Artifacts: v1 BC/RL policies for all 4 (drive × world) combos trained
in clean gym (success: diff/field 90-100%, diff/round 58%, mec/field
60-100%, mec/round 50-100%). DAgger r1/r2 BCs for diff/field show
12%→38% progression on gym HERDING_WEBOTS proxy but did not close
to actual Webots LiDAR (0/5 throughout). Next: LSTM policy or
learned tracker per the project-state memory.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
2026-05-16 17:21:02 +00:00

275 lines
12 KiB
Python

"""Collect (obs, action) demonstrations from an analytic teacher.
Runs the chosen teacher across a grid of ``(n_sheep, seed)`` combos at
full difficulty, logs every Nth ``(obs, action)`` pair, and saves
successful trajectories to ``.npz`` for behaviour cloning. The teacher
is wrapped in :class:`ActiveScanTeacher` by default so it operates on
the same partial-obs view the student will have at deployment.
Usage::
python -m training.bc.collect --teacher strombom \\
--out training/bc/demos.npz --frame-stack 4
"""
from __future__ import annotations
import argparse
import os
import time
from pathlib import Path
import numpy as np
# Configure field geometry before other herding imports read it at module level.
from herding.world.geometry import configure_from_args as _configure_from_args
_configure_from_args()
from herding.control.active_scan import ActiveScanTeacher
from herding.world.geometry import PEN_ENTRY, FIELD_SHAPE
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 training.herding_env import HerdingEnv
TEACHERS = {
"sequential": sequential_action,
"strombom": strombom_action,
"universal": universal_action,
}
def _call_teacher(fn, dog_xy, dog_heading, sheep_positions, pen_target,
drive_mode="differential"):
"""Call any teacher function and return (vx, vy, omega, mode).
Normalizes across 3-tuple teachers (vx, vy, mode) and 4-tuple
universal teacher (vx, vy, omega, mode). ActiveScanTeacher (when
invoked with drive_mode="mecanum") propagates the base teacher's
omega — see test_active_scan_preserves_mecanum_omega.
"""
# The universal teacher and ActiveScanTeacher accept the extended
# (dog_xy, heading, sheep, pen, drive_mode) signature. Older
# teachers accept (dog_xy, sheep, pen). Detect by trying the
# extended call first.
try:
result = fn(dog_xy, dog_heading, sheep_positions, pen_target,
drive_mode)
except TypeError:
try:
result = fn(dog_xy, dog_heading, sheep_positions, pen_target)
except TypeError:
result = fn(dog_xy, sheep_positions, pen_target)
if len(result) == 4:
return result # (vx, vy, omega, mode)
vx, vy, mode = result
return vx, vy, 0.0, mode
def collect_one(n_sheep: int, seed: int, max_steps: int, subsample: int,
teacher_fn, frame_stack: int = 1, privileged: bool = False,
drive_mode: str = "differential", herding_cfg=None,
actor_policy=None):
"""Collect (obs, teacher_action) pairs from one episode.
``actor_policy`` (DAgger mode): a callable ``policy(obs) -> action`` that
drives the env. The teacher still labels each visited state. If ``None``
(default), the teacher drives.
"""
env = HerdingEnv(n_sheep=n_sheep, max_steps=max_steps,
difficulty=1.0, seed=seed, frame_stack=frame_stack,
drive_mode=drive_mode, herding_cfg=herding_cfg)
obs, _ = env.reset(seed=seed)
obs_list, action_list = [], []
scan_teacher = ActiveScanTeacher(teacher_fn)
for step in range(max_steps):
if privileged:
positions = {f"s{i}": (float(env.sheep_x[i]), float(env.sheep_y[i]))
for i in range(env.n_sheep) if not env.sheep_penned[i]}
if not positions:
break
vx, vy, omega, _mode = _call_teacher(
teacher_fn, (env.dog_x, env.dog_y), env.dog_heading,
positions, PEN_ENTRY, drive_mode,
)
else:
positions = env.perceived_positions()
result = _call_teacher(
scan_teacher, (env.dog_x, env.dog_y), env.dog_heading,
positions, PEN_ENTRY, drive_mode,
)
vx, vy, omega, _mode = result
if drive_mode == "mecanum":
teacher_action = np.array([vx, vy, omega], dtype=np.float32)
else:
teacher_action = np.array([vx, vy], dtype=np.float32)
if step % subsample == 0:
obs_list.append(obs.copy())
action_list.append(teacher_action.copy())
# In DAgger mode the policy drives; otherwise the teacher does.
step_action = (actor_policy(obs) if actor_policy is not None
else teacher_action)
obs, _r, term, trunc, _info = env.step(step_action)
if term or trunc:
break
success = bool(env.sheep_penned.all())
return (
np.asarray(obs_list, dtype=np.float32),
np.asarray(action_list, dtype=np.float32),
success,
env.steps,
)
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--out", default="training/bc/demos.npz")
parser.add_argument("--n-sheep-list", default="1,2,3,5,8,10")
parser.add_argument("--seeds-per-n", type=int, default=15)
parser.add_argument("--max-steps", type=int, default=30000)
parser.add_argument("--subsample", type=int, default=5,
help="Keep every Nth (obs, action) pair.")
parser.add_argument("--keep-failures", action="store_true",
help="Include partial-success trajectories. Default off.")
parser.add_argument("--teacher", default="universal",
choices=list(TEACHERS.keys()),
help="Which analytic teacher to demonstrate.")
parser.add_argument("--frame-stack", type=int, default=1,
help="Concatenate the last K obs into a "
"(32·K)-D vector for the policy.")
parser.add_argument("--privileged", action="store_true",
help="Teacher reads ground truth instead of "
"tracker output (asymmetric BC).")
parser.add_argument("--drive-mode", default="differential",
choices=["differential", "mecanum"],
help="Drive mode for the dog robot.")
parser.add_argument("--world", default=None,
choices=["field", "field_round"],
help="World shape. If not set, uses HERDING_WORLD "
"env var or defaults to 'field'. Must be set "
"before geometry is imported.")
# Domain randomisation — applied to the gym env during collection so
# the teacher demonstrates under the same noise the policy will face.
parser.add_argument("--fp-rate", type=float, default=0.0,
help="Mean false-positive detections injected per "
"step (Poisson λ). 0 = clean sim (default).")
parser.add_argument("--action-smooth", type=float, default=0.0,
help="EMA coefficient on dog actions (0 = none). "
"Set to 0.55 to match the Webots controller.")
parser.add_argument("--wheel-slip-std", type=float, default=0.0,
help="Gaussian noise (rad/s) on wheel speeds for "
"mecanum dynamics domain randomisation.")
parser.add_argument("--dagger-policy", default=None,
help="Path to a BC/PPO policy directory. When set, "
"the policy drives the env (DAgger) while the "
"teacher labels every visited state.")
parser.add_argument("--use-webots-preset", action="store_true",
help="Use HERDING_WEBOTS preset (140° FOV + tight "
"tracker). Match this to deployment for DAgger.")
args = parser.parse_args()
# Validate --world matches geometry (already configured by the
# early pre-parse above, or by HERDING_WORLD env var).
if args.world is not None and args.world != FIELD_SHAPE:
print(f"[demos] WARNING: --world={args.world} but geometry is "
f"'{FIELD_SHAPE}'. This should not happen — file a bug.")
from herding.config import HerdingConfig, HERDING_WEBOTS, DomainRandomConfig, RobotConfig
if args.use_webots_preset:
herding_cfg = HERDING_WEBOTS.replace(
domain_random=DomainRandomConfig(
fp_rate=args.fp_rate,
wheel_slip_std=args.wheel_slip_std,
),
robot=RobotConfig(action_smooth=args.action_smooth),
)
print(f"[demos] HERDING_WEBOTS preset + DR: fp_rate={args.fp_rate} "
f"action_smooth={args.action_smooth} wheel_slip_std={args.wheel_slip_std}")
else:
herding_cfg = None
if args.fp_rate > 0.0 or args.action_smooth > 0.0 or args.wheel_slip_std > 0.0:
herding_cfg = HerdingConfig(
domain_random=DomainRandomConfig(
fp_rate=args.fp_rate,
wheel_slip_std=args.wheel_slip_std,
),
robot=RobotConfig(action_smooth=args.action_smooth),
)
print(f"[demos] domain-random: fp_rate={args.fp_rate} "
f"action_smooth={args.action_smooth} "
f"wheel_slip_std={args.wheel_slip_std}")
actor_policy = None
if args.dagger_policy is not None:
# DAgger: failures are the most valuable data (off-policy states
# where the student needs teacher correction). Always keep them.
args.keep_failures = True
from stable_baselines3 import PPO
from pathlib import Path as _P
run = _P(args.dagger_policy)
for name in ("policy.zip", "final.zip"):
if (run / name).exists():
zip_path = run / name
break
else:
raise FileNotFoundError(
f"No policy found in {run} (tried policy.zip, final.zip)")
_model = PPO.load(str(zip_path), device="auto")
print(f"[demos] DAgger mode: actor = {zip_path}")
def actor_policy(obs):
obs_b = np.asarray(obs, dtype=np.float32).reshape(1, -1)
a, _ = _model.predict(obs_b, deterministic=True)
return a[0]
teacher_fn = TEACHERS[args.teacher]
print(f"[demos] teacher: {args.teacher} world: {FIELD_SHAPE}")
n_sheep_list = [int(x) for x in args.n_sheep_list.split(",")]
print(f"[demos] grid: n_sheep={n_sheep_list}, seeds={args.seeds_per_n}, "
f"max_steps={args.max_steps}, subsample={args.subsample}")
all_obs, all_actions, all_meta = [], [], []
t_start = time.time()
n_success = 0; n_total = 0
for n in n_sheep_list:
for seed in range(args.seeds_per_n):
obs, actions, success, total_steps = collect_one(
n, seed, args.max_steps, args.subsample, teacher_fn,
frame_stack=args.frame_stack, privileged=args.privileged,
drive_mode=args.drive_mode, herding_cfg=herding_cfg,
actor_policy=actor_policy,
)
n_total += 1
if success:
n_success += 1
keep = success or args.keep_failures
if keep and len(obs) > 0:
all_obs.append(obs)
all_actions.append(actions)
all_meta.append((n, seed, len(obs), int(success), total_steps))
tag = "✓" if success else "✗"
print(f" [{tag}] n={n:>2d} seed={seed:>2d} steps={total_steps:>6d} "
f"logged={len(obs):>5d}")
if not all_obs:
raise RuntimeError("No trajectories kept — try --keep-failures.")
obs = np.concatenate(all_obs, axis=0)
actions = np.concatenate(all_actions, axis=0)
meta = np.array(all_meta, dtype=np.int32)
Path(args.out).parent.mkdir(parents=True, exist_ok=True)
np.savez(args.out, obs=obs, actions=actions, meta=meta)
elapsed = time.time() - t_start
print(f"\n=== {n_success}/{n_total} trajectories successful ({100*n_success/n_total:.0f}%) ===")
print(f"=== {len(obs)} transitions saved to {args.out} ===")
print(f"=== obs={obs.shape}, actions={actions.shape}, elapsed={elapsed:.0f}s ===")
if __name__ == "__main__":
main()