jalf/TIR_PROJ
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Checkpoint 6

Browse Source
This commit is contained in:
Johnny Fernandes
2026-05-11 10:35:48 +01:00
parent b457155538
commit fce0e0c786
27 changed files with 194 additions and 704 deletions
Download Patch File Download Diff File Expand all files Collapse all files
training/README.md
+35 -54
View File
@@ -1,21 +1,16 @@
# Training pipeline
Behavior cloning of analytic herding teachers into a neural-network
policy that runs under LiDAR perception in Webots.
Two stages, strictly sequential:
```
sim demos (active-scan teacher on tracker output, K=4 frame stack)
sim demos (Strömbom on tracker output, K=4 frame stack)
│
▼
bc_pretrain.py ──► runs/bc (BC baseline)
bc_pretrain.py ──► runs/bc (Strömbom-imitated MLP)
│
▼ KL-regularised PPO fine-tune (training/train_ppo.py)
▼ KL-regularised PPO fine-tune
│
runs/rl (deployed `rl` mode)
# optional branch — kept for reference, not deployed:
runs/bc_dagger (Webots-grounded DAgger refinement, useful if a
modified world breaks sim-to-real transfer)
runs/rl (deployed `rl` mode — beats BC and Strömbom)
```
## Files
@@ -23,10 +18,9 @@ runs/bc_dagger (Webots-grounded DAgger refinement, useful if a
```
herding_env.py — Gymnasium env (LiDAR raycast + tracker by default)
bc_pretrain.py — MSE + cosine BC of (obs, action) demos into MlpPolicy
eval.py — analytic teachers + BC policies, full n=1..10 grid
parity_test.py — shape / determinism / baseline smoke test
runs/ — checkpoints (most are .gitignored; the deployed
ones are whitelisted in the top-level .gitignore)
train_ppo.py — KL-regularised PPO fine-tune of a BC checkpoint
eval.py — multi-seed analytic / learned policy comparison
runs/ — checkpoints (whitelisted entries in top-level .gitignore)
```
## Setup
@@ -39,75 +33,62 @@ CPU is the default and recommended device — SB3 PPO with an MLP policy
of this size runs faster on CPU than GPU because the bottleneck is
rollout collection, not gradient compute.
## The BC pipeline
## End-to-end pipeline
```
```bash
# 1. Sim demos with the active-scan + Strömbom teacher under LiDAR
# perception. K=4 frame stack so the MLP has temporal context.
python -m tools.collect_demos --teacher strombom \
--out demos.npz --seeds-per-n 15 --subsample 3 --frame-stack 4
--out training/demos.npz --seeds-per-n 15 --subsample 3 --frame-stack 4
# 2. Behavior-clone.
python -m training.bc_pretrain --demos demos.npz \
--out runs/bc --epochs 60 --net-arch 512,512
# 2. Behaviour-clone.
python -m training.bc_pretrain --demos training/demos.npz \
--out training/runs/bc --epochs 60 --net-arch 512,512
# 3. Evaluate.
python -m training.eval --policy runs/bc \
--max-flock 10 --max-steps 8000 --n-seeds 5
# 3. KL-regularised PPO fine-tune of bc.
python -m training.train_ppo \
--bc training/runs/bc --out training/runs/rl \
--total-timesteps 1000000
# 4. Multi-seed eval (env-side, fast).
python -m training.eval --policy training/runs/rl \
--max-flock 10 --max-steps 15000 --n-seeds 10
```
`bc_pretrain.py` saves the **best-val_cos** snapshot, not the final
epoch — multi-modal teachers make training noisy and the last epoch is
often worse than an earlier one.
## DAgger from Webots
Sim-only BC plateaus because the env's 2D raycast can't reproduce all
the false-positive clusters Webots generates from real geometry. The
fix is to collect (obs, teacher_action) pairs from inside Webots:
```
# Headless DAgger collection: 5 flock sizes × 3 runs each.
tools/auto_dagger.sh 3 60
# Merge with the sim baseline + retrain.
python -m tools.dagger_merge_train --out runs/bc_dagger
```
Iterate by re-running collection with the new student in the driver's
seat:
```
HERDING_POLICY_DIR=$PWD/training/runs/bc_dagger \
HERDING_DAGGER_DRIVER=student \
tools/auto_dagger.sh 3 60
python -m tools.dagger_merge_train --out runs/bc_dagger
```
`train_ppo.py` loads BC weights into both a trainable policy and a
frozen reference, fixes `log_std` small, and adds `β · KL(π‖π_ref)` to
the loss so the policy can only move within a trust region around BC.
See the file header for hyperparameter rationale.
## Available analytic teachers
| Name | What it does | Notes |
|---|---|---|
| `strombom` | Canonical Strömbom — collect when flock is scattered, drive CoM otherwise | Default; works well for n=1–10 under tight cohesion |
| `strombom` | Strömbom 2014 — collect when flock is scattered, drive CoM otherwise | Default; works for n=1–10 under tight cohesion |
| `sequential` | Pick the sheep closest to the pen and drive only it | Alternative; needs loose-cohesion regime |
Both are wrapped at demo-collection time in
`herding/active_scan.py:ActiveScanTeacher`, which adds an opening
in-place rotation, walk-to-centre when the LiDAR sees nothing, and
near-sheep speed modulation (the same modulation `herding/control.py`
applies to every dog mode at inference).
`herding/control/active_scan.py:ActiveScanTeacher`, which adds an
opening in-place rotation, walk-to-centre when the LiDAR sees
nothing, and near-sheep speed modulation (same modulation
`herding/control/modulation.py` applies to every dog mode at
inference).
## Evaluating analytic teachers directly
```
python -m training.eval --policy strombom --max-flock 10 --max-steps 8000 --n-seeds 5
python -m training.eval --policy sequential --max-flock 10 --max-steps 8000 --n-seeds 5
python -m training.eval --policy strombom --max-flock 10 --max-steps 15000 --n-seeds 10
python -m training.eval --policy sequential --max-flock 10 --max-steps 15000 --n-seeds 10
```
## Webots inference
```
tools/run_webots.sh 10 rl
tools/run_webots.sh 10 bc # or rl, strombom, sequential
```
The dog controller loads `runs/bc` for `bc` mode and `runs/rl` for
training/eval.py
+3 -3
View File
@@ -25,9 +25,9 @@ if _PROJECT_ROOT not in sys.path:
import numpy as np
from herding.geometry import MAX_SHEEP, PEN_ENTRY
from herding.sequential import compute_action as sequential_action
from herding.strombom import compute_action as strombom_action
from herding.world.geometry import MAX_SHEEP, PEN_ENTRY
from herding.control.sequential import compute_action as sequential_action
from herding.control.strombom import compute_action as strombom_action
from training.herding_env import HerdingEnv
training/herding_env.py
+7 -7
View File
@@ -56,24 +56,24 @@ _PROJECT_ROOT = os.path.normpath(os.path.join(_HERE, ".."))
if _PROJECT_ROOT not in sys.path:
sys.path.insert(0, _PROJECT_ROOT)
from herding.diffdrive import (
from herding.world.diffdrive import (
heading_speed_to_wheels, kinematics_step, velocity_to_wheels,
)
from herding.flocking_sim import (
from herding.world.flocking_sim import (
FLEE_SPEED, MAX_SPEED, WANDER_SPEED, compute_heading_speed,
)
from herding.geometry import (
from herding.world.geometry import (
DOG_MAX_LINEAR, DOG_MAX_WHEEL_OMEGA, DOG_SOUTH_LIMIT, DOG_WHEEL_BASE,
DOG_WHEEL_RADIUS, FIELD_X, FIELD_Y, GATE_X, MAX_SHEEP,
PEN_ENTRY, PEN_X, PEN_Y,
SHEEP_MAX_WHEEL_OMEGA, SHEEP_WHEEL_BASE, SHEEP_WHEEL_RADIUS,
WEBOTS_DT, is_penned_position,
)
from herding.lidar_perception import detections_from_scan
from herding.lidar_sim import simulate_scan
from herding.perception.lidar_perception import detections_from_scan
from herding.perception.lidar_sim import simulate_scan
from herding.obs import OBS_DIM, build_obs
from herding.sheep_tracker import SheepTracker
from herding.strombom import compute_action as strombom_action
from herding.perception.sheep_tracker import SheepTracker
from herding.control.strombom import compute_action as strombom_action
class HerdingEnv(gym.Env):
training/parity_test.py
-96
View File
@@ -1,96 +0,0 @@
"""Parity smoke-test for the herding env.
Verifies (a) all imports resolve, (b) the env's reset/step contract is
correct, (c) deterministic seeds give deterministic trajectories, and
(d) the Strömbom baseline can drive the env without crashing.
Run::
python -m training.parity_test
"""
from __future__ import annotations
import os
import sys
_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)
import numpy as np
from herding.geometry import MAX_SHEEP, PEN_ENTRY
from herding.obs import OBS_DIM
from herding.strombom import compute_action
from training.herding_env import HerdingEnv
def test_obs_action_shapes():
env = HerdingEnv(n_sheep=3, seed=0)
obs, info = env.reset()
assert obs.shape == (OBS_DIM,), obs.shape
assert obs.dtype == np.float32
obs2, r, term, trunc, info = env.step(np.array([0.5, 0.0], dtype=np.float32))
assert obs2.shape == (OBS_DIM,)
assert isinstance(r, float)
assert isinstance(term, bool) and isinstance(trunc, bool)
print("[ok] shapes")
def test_reset_determinism():
"""Reset with the same seed should give the same initial observation.
We don't require step-determinism — PPO doesn't need it, and chasing
bit-exactness through the flocking jitter isn't worth the complexity.
"""
env_a = HerdingEnv(n_sheep=3, seed=42)
env_b = HerdingEnv(n_sheep=3, seed=42)
obs_a, _ = env_a.reset(seed=42)
obs_b, _ = env_b.reset(seed=42)
assert np.allclose(obs_a, obs_b), "Reset is non-deterministic for same seed"
print("[ok] reset determinism")
def test_curriculum_n_sheep_varies():
env = HerdingEnv(seed=0)
sizes = set()
for _ in range(40):
_, info = env.reset()
sizes.add(info["n_sheep"])
assert 1 in sizes
assert max(sizes) <= MAX_SHEEP
print(f"[ok] curriculum sampling — saw n_sheep in {sorted(sizes)}")
def test_strombom_drives_env():
"""Quick functional check that the analytic baseline can play the env
without exploding. Not a success-rate test — just no errors / NaNs."""
env = HerdingEnv(n_sheep=2, max_steps=400, seed=1)
obs, _ = env.reset()
for t in range(400):
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, _mode = compute_action((env.dog_x, env.dog_y), positions, PEN_ENTRY)
obs, r, term, trunc, info = env.step(np.array([vx, vy], dtype=np.float32))
assert np.isfinite(obs).all(), f"NaN/Inf in obs at step {t}"
assert np.isfinite(r), f"NaN reward at step {t}"
if term or trunc:
break
print(f"[ok] strombom rollout — final n_penned={int(env.sheep_penned.sum())}/{env.n_sheep} after {env.steps} steps")
def main():
test_obs_action_shapes()
test_reset_determinism()
test_curriculum_n_sheep_varies()
test_strombom_drives_env()
print("\nAll parity checks passed.")
if __name__ == "__main__":
main()