From 7ab69ab0f336d2f454b35929488fd25cbff35f20 Mon Sep 17 00:00:00 2001 From: Johnny Fernandes Date: Sun, 17 May 2026 01:58:15 +0000 Subject: [PATCH] Rename multi-segment functions to two-concept names; polish docstrings MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Naming pass: rename functions whose third+ segment is redundant or implementation-detail, sticking to the codebase's preferred ``noun_verb`` / ``verb_noun`` two-concept idiom. Renames are atomic across definitions, callers, and tests. is_penned_position → is_penned modulate_speed_near_sheep → modulate_speed mecanum_kinematics_step → mecanum_step policy_forward_mean → forward_mean Two-concept patterns like ``velocity_to_wheels`` / ``detections_from_scan`` / ``make_strombom_predictor`` are left alone — they're idiomatic converters / factories that read as a single concept, and the longer form aids grep-ability. Docstring polish: * ``herding/config.py`` header drops the "previously lived as a module-level literal" historical framing — we ship as a single thing, so the refactor anecdote no longer earns its keep. The usage examples now mention both ``HERDING_WEBOTS`` and ``HERDING_MEC_WEBOTS`` presets. 126 pytest cases still pass. Co-Authored-By: Claude Opus 4.7 --- controllers/sheep/sheep.py | 4 +-- controllers/shepherd_dog/shepherd_dog.py | 16 +++++------ herding/config.py | 34 +++++++++++------------- herding/control/active_scan.py | 6 ++--- herding/control/modulation.py | 4 +-- herding/perception/sheep_tracker.py | 8 +++--- herding/world/diffdrive.py | 4 +-- herding/world/geometry.py | 2 +- tests/test_control.py | 18 ++++++------- tests/test_diffdrive.py | 16 +++++------ tests/test_geometry.py | 24 ++++++++--------- tests/test_perception.py | 2 +- training/bc/pretrain.py | 6 ++--- training/herding_env.py | 8 +++--- 14 files changed, 75 insertions(+), 77 deletions(-) diff --git a/controllers/sheep/sheep.py b/controllers/sheep/sheep.py index a6fdd0d..a319836 100644 --- a/controllers/sheep/sheep.py +++ b/controllers/sheep/sheep.py @@ -27,7 +27,7 @@ from herding.world.diffdrive import heading_speed_to_wheels from herding.world.flocking_sim import MAX_SPEED, compute_heading_speed from herding.world.geometry import ( SHEEP_MAX_WHEEL_OMEGA, - is_penned_position, + is_penned, ) @@ -92,7 +92,7 @@ while robot.step(timestep) != -1: pos = gps.getValues() x, y = pos[0], pos[1] - if not penned and is_penned_position(x, y): + if not penned and is_penned(x, y): penned = True paint_pink() diff --git a/controllers/shepherd_dog/shepherd_dog.py b/controllers/shepherd_dog/shepherd_dog.py index 1af4ed3..97ae6b2 100644 --- a/controllers/shepherd_dog/shepherd_dog.py +++ b/controllers/shepherd_dog/shepherd_dog.py @@ -85,7 +85,7 @@ import numpy as np from controller import Robot from herding.control.active_scan import ActiveScanTeacher -from herding.control.modulation import modulate_speed_near_sheep +from herding.control.modulation import modulate_speed 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 @@ -95,7 +95,7 @@ from herding.perception.sheep_tracker import SheepTracker from herding.world.diffdrive import velocity_to_mecanum_wheels, velocity_to_wheels from herding.world.geometry import ( DOG_SOUTH_LIMIT, - PEN_ENTRY, is_penned_position, + PEN_ENTRY, is_penned, ) from herding.config import HERDING_WEBOTS, RobotConfig @@ -332,7 +332,7 @@ if MODE == "calibrate": _pos0 = gps.getValues(); _x0, _y0 = _pos0[0], _pos0[1] _n_calib = compass.getValues(); _h0 = math.atan2(_n_calib[0], _n_calib[1]) # Gym-predicted displacement using shared kinematics. - from herding.world.diffdrive import velocity_to_mecanum_wheels, mecanum_kinematics_step + from herding.world.diffdrive import velocity_to_mecanum_wheels, mecanum_step from herding.world.geometry import WEBOTS_DT as _DT _xg, _yg, _hg = _x0, _y0, _h0 for _ in range(_calib_n): @@ -343,7 +343,7 @@ if MODE == "calibrate": max_wheel_omega=DOG_MAX_WHEEL_OMEGA, k_turn=4.0, wheel_base=DOG_WHEEL_BASE, ) - _xg, _yg, _hg = mecanum_kinematics_step( + _xg, _yg, _hg = mecanum_step( _xg, _yg, _hg, _wfl, _wfr, _wrl, _wrr, DOG_WHEEL_RADIUS, DOG_WHEEL_BASE_X / 2, DOG_WHEEL_BASE_Y / 2, _DT, ) @@ -414,7 +414,7 @@ while robot.step(timestep) != -1: if USE_GT_PERCEPTION and _gt_sheep: # Bypass tracker: feed GT emitter positions directly to policy/teacher. sheep_positions = {k: v for k, v in _gt_sheep.items() - if not is_penned_position(v[0], v[1])} + if not is_penned(v[0], v[1])} tracker.update(detections) # still advance tracker for diagnostics else: sheep_positions = tracker.update(detections) @@ -453,7 +453,7 @@ while robot.step(timestep) != -1: vx, vy, _mode_str = result # Near-sheep speed modulation (shared by every mode). - vx, vy = modulate_speed_near_sheep(vx, vy, dog_xy, sheep_positions) + vx, vy = modulate_speed(vx, vy, dog_xy, sheep_positions) # EMA smoothing — kills frame-to-frame action jitter. if DRIVE_MODE == "mecanum": @@ -489,7 +489,7 @@ while robot.step(timestep) != -1: # fire during the first few steps while the receiver fills. if _gt_sheep and not _run_done: gt_active = sum(1 for x, y in _gt_sheep.values() - if not is_penned_position(x, y)) + if not is_penned(x, y)) if gt_active == 0: os.makedirs(os.path.dirname(RUN_DONE_FILE), exist_ok=True) open(RUN_DONE_FILE, "w").close() @@ -499,7 +499,7 @@ while robot.step(timestep) != -1: if step_count % 200 == 0: gt_penned = sum(1 for x, y in _gt_sheep.values() - if is_penned_position(x, y)) + if is_penned(x, y)) gt_total = len(_gt_sheep) common = (f"[dog mode={MODE} drive={DRIVE_MODE}] step={step_count} " f"GT_penned={gt_penned}/{gt_total} " diff --git a/herding/config.py b/herding/config.py index bdf2feb..1d99de6 100644 --- a/herding/config.py +++ b/herding/config.py @@ -1,32 +1,30 @@ """Central configuration dataclasses for the herding simulation. -Every tunable constant that previously lived as a module-level literal in -perception/lidar_sim.py, perception/lidar_perception.py, -perception/sheep_tracker.py, world/geometry.py, or training/herding_env.py -is now represented here as a field with its original default value. +Every tunable parameter lives here as a frozen dataclass field — LiDAR +spec, cluster detection thresholds, tracker gates, robot kinematics, +and domain-randomisation knobs — composed into :class:`HerdingConfig`. -Usage — use the module defaults unchanged:: +Usage — accept the defaults:: - env = HerdingEnv() # same behaviour as before + env = HerdingEnv() -Override a subset of parameters:: +Override a subset:: - from herding.config import HerdingConfig, TrackerConfig cfg = HerdingConfig(tracker=TrackerConfig(forget_steps=60)) env = HerdingEnv(herding_cfg=cfg) -Use a named preset for Webots-matched training:: +Use a named preset:: - from herding.config import HERDING_WEBOTS - env = HerdingEnv(herding_cfg=HERDING_WEBOTS) + env = HerdingEnv(herding_cfg=HERDING_WEBOTS) # 140° FOV + env = HerdingEnv(herding_cfg=HERDING_MEC_WEBOTS) # + mecanum slip Design notes ------------ -* All dataclasses are frozen — instances are immutable after construction. -* This module must not import from other ``herding.*`` packages to avoid - import cycles. Field-geometry constants (pen coordinates, field size) - stay in ``herding.world.geometry`` because they depend on the world - variant selected at runtime via ``HERDING_WORLD``. +* All dataclasses are frozen so instances are immutable after construction. +* This module must not import from other ``herding.*`` packages — + field-geometry constants live in ``herding.world.geometry`` because + they depend on the world variant selected at runtime via + ``HERDING_WORLD``, which would create an import cycle here. """ from __future__ import annotations @@ -257,7 +255,7 @@ class RobotConfig: ``1.0`` (default) = perfect mecanum kinematics. ``0.4`` matches the Webots roller-hinge mecanum proto calibration (62% slip on strafe, - 11% on forward). Used by ``mecanum_kinematics_step`` only — has no + 11% on forward). Used by ``mecanum_step`` only — has no effect on differential drive. """ @@ -266,7 +264,7 @@ class RobotConfig: ``0.0`` (default) = no bleed. ``-0.28`` matches the Webots proto's consistent backward push under strafe commands. Used by - ``mecanum_kinematics_step`` only. + ``mecanum_step`` only. """ def __post_init__(self) -> None: diff --git a/herding/control/active_scan.py b/herding/control/active_scan.py index df7f815..b4130e9 100644 --- a/herding/control/active_scan.py +++ b/herding/control/active_scan.py @@ -12,7 +12,7 @@ exploration behaviours: beyond the 12 m LiDAR range). When the tracker has detections the base teacher's action is used, -post-processed by ``modulate_speed_near_sheep`` so the dog doesn't +post-processed by ``modulate_speed`` so the dog doesn't charge the flock. """ @@ -20,7 +20,7 @@ from __future__ import annotations import math -from herding.control.modulation import modulate_speed_near_sheep +from herding.control.modulation import modulate_speed INITIAL_SCAN_STEPS = 80 # ≈1.3 s — covers one full rotation @@ -117,6 +117,6 @@ class ActiveScanTeacher: else: vx, vy, mode = result omega = 0.0 - vx, vy = modulate_speed_near_sheep(vx, vy, dog_xy, sheep_positions) + vx, vy = modulate_speed(vx, vy, dog_xy, sheep_positions) self.last_action = (vx, vy) return vx, vy, omega, mode diff --git a/herding/control/modulation.py b/herding/control/modulation.py index 55cd9b5..e6ddb09 100644 --- a/herding/control/modulation.py +++ b/herding/control/modulation.py @@ -1,6 +1,6 @@ """Shared action post-processing. -Every dog mode routes its action through ``modulate_speed_near_sheep`` +Every dog mode routes its action through ``modulate_speed`` so the magnitude is reduced near sheep — direction (intent) is preserved. """ @@ -14,7 +14,7 @@ SLOW_NEAR_SHEEP = 2.5 # m — distance below which action norm is scaled down MIN_SPEED = 0.30 # action norm at zero distance -def modulate_speed_near_sheep( +def modulate_speed( vx: float, vy: float, dog_xy: tuple[float, float], sheep_positions, diff --git a/herding/perception/sheep_tracker.py b/herding/perception/sheep_tracker.py index 8e0f0a2..f4bf023 100644 --- a/herding/perception/sheep_tracker.py +++ b/herding/perception/sheep_tracker.py @@ -15,7 +15,7 @@ Three-stage greedy nearest-neighbour data association: extrapolated for up to ``PREDICT_STEPS`` frames, then falls back to last-seen static memory until ``FORGET_STEPS`` deletes it. 3. **Pen latching**. A track whose estimated position crosses the gate - plane south of ``is_penned_position`` is marked penned, excluded + plane south of ``is_penned`` is marked penned, excluded from ``get_positions``, and kept indefinitely. Output of :meth:`SheepTracker.get_positions` is ``{name: (x, y)}`` — @@ -31,7 +31,7 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: from herding.config import TrackerConfig -from herding.world.geometry import MAX_SHEEP, in_pen, is_penned_position +from herding.world.geometry import MAX_SHEEP, in_pen, is_penned GATE_M = 2.5 # m — primary NN gate (recently observed tracks) @@ -356,10 +356,10 @@ class SheepTracker: at y ≈ -15) from being permanently latched as penned tracks. """ from herding.world.geometry import GATE_Y - # Apply depth threshold to both in_pen and is_penned_position so + # Apply depth threshold to both in_pen and is_penned so # that any position in the gate column must clear GATE_Y - depth. threshold = GATE_Y - self._pen_latch_depth - return (in_pen(x, y) or is_penned_position(x, y)) and y <= threshold + return (in_pen(x, y) or is_penned(x, y)) and y <= threshold def get_positions(self, min_freshness: int | None = None) -> dict[str, tuple[float, float]]: """Promoted (non-candidate, non-penned) tracks as ``{name: (x, y)}``. diff --git a/herding/world/diffdrive.py b/herding/world/diffdrive.py index 6f4e8d8..88da44f 100644 --- a/herding/world/diffdrive.py +++ b/herding/world/diffdrive.py @@ -3,7 +3,7 @@ controllers. First-order rigid-body model — no slip, wheel-accel limits, or contact forces by default. Pass ``slip_std`` and an ``rng`` to -:func:`kinematics_step` / :func:`mecanum_kinematics_step` to add +:func:`kinematics_step` / :func:`mecanum_step` to add per-wheel Gaussian speed noise for domain randomisation. """ @@ -80,7 +80,7 @@ def heading_speed_to_wheels(heading, speed_motor, h, max_wheel_omega, # Mecanum (4-wheel omnidirectional) kinematics # --------------------------------------------------------------------------- -def mecanum_kinematics_step(x, y, h, w_fl, w_fr, w_rl, w_rr, +def mecanum_step(x, y, h, w_fl, w_fr, w_rl, w_rr, wheel_radius, lx, ly, dt, slip_std: float = 0.0, rng: Optional[np.random.Generator] = None, diff --git a/herding/world/geometry.py b/herding/world/geometry.py index f21aedb..0c9243b 100644 --- a/herding/world/geometry.py +++ b/herding/world/geometry.py @@ -171,7 +171,7 @@ def in_gate_corridor(x: float, y: float, margin: float = 0.0) -> bool: and PEN_Y[0] - margin <= y <= GATE_Y + margin) -def is_penned_position(x: float, y: float, latch_margin: float = 0.2) -> bool: +def is_penned(x: float, y: float, latch_margin: float = 0.2) -> bool: """True iff (x, y) is in the gate column and south of the gate line.""" return (GATE_X[0] - latch_margin <= x <= GATE_X[1] + latch_margin and y <= GATE_Y) diff --git a/tests/test_control.py b/tests/test_control.py index 429dfc1..cf12d5c 100644 --- a/tests/test_control.py +++ b/tests/test_control.py @@ -8,7 +8,7 @@ from herding.control.active_scan import ( EMPTY_DEBOUNCE_STEPS, INITIAL_SCAN_STEPS, ActiveScanTeacher, ) from herding.control.modulation import ( - MIN_SPEED, SLOW_NEAR_SHEEP, modulate_speed_near_sheep, + MIN_SPEED, SLOW_NEAR_SHEEP, modulate_speed, ) from herding.control.sequential import compute_action as sequential_action from herding.control.strombom import ( @@ -23,23 +23,23 @@ from herding.world.geometry import PEN_ENTRY # --------------------------------------------------------------------------- def test_modulation_empty_input_passthrough(): - assert modulate_speed_near_sheep(1.0, 0.0, (0.0, 0.0), []) == (1.0, 0.0) - assert modulate_speed_near_sheep(1.0, 0.0, (0.0, 0.0), {}) == (1.0, 0.0) + assert modulate_speed(1.0, 0.0, (0.0, 0.0), []) == (1.0, 0.0) + assert modulate_speed(1.0, 0.0, (0.0, 0.0), {}) == (1.0, 0.0) def test_modulation_far_sheep_passthrough(): - vx, vy = modulate_speed_near_sheep(1.0, 0.0, (0.0, 0.0), [(100.0, 0.0)]) + vx, vy = modulate_speed(1.0, 0.0, (0.0, 0.0), [(100.0, 0.0)]) assert (vx, vy) == (1.0, 0.0) def test_modulation_close_sheep_min_speed(): - vx, vy = modulate_speed_near_sheep(1.0, 0.0, (0.0, 0.0), [(0.0, 0.0)]) + vx, vy = modulate_speed(1.0, 0.0, (0.0, 0.0), [(0.0, 0.0)]) assert math.isclose(vx, MIN_SPEED) assert vy == 0.0 def test_modulation_preserves_direction(): - vx, vy = modulate_speed_near_sheep(0.6, 0.8, (0.0, 0.0), [(1.0, 0.0)]) + vx, vy = modulate_speed(0.6, 0.8, (0.0, 0.0), [(1.0, 0.0)]) ratio = math.hypot(vx, vy) # Direction preserved. assert math.isclose(vx / ratio, 0.6, abs_tol=1e-6) @@ -47,16 +47,16 @@ def test_modulation_preserves_direction(): def test_modulation_linear_ramp_midpoint(): - vx, _ = modulate_speed_near_sheep(1.0, 0.0, (0.0, 0.0), + vx, _ = modulate_speed(1.0, 0.0, (0.0, 0.0), [(SLOW_NEAR_SHEEP / 2, 0.0)]) expected = MIN_SPEED + (1.0 - MIN_SPEED) * 0.5 assert math.isclose(vx, expected, abs_tol=1e-6) def test_modulation_accepts_dict_input(): - vx_list, _ = modulate_speed_near_sheep(1.0, 0.0, (0.0, 0.0), + vx_list, _ = modulate_speed(1.0, 0.0, (0.0, 0.0), [(1.0, 0.0)]) - vx_dict, _ = modulate_speed_near_sheep(1.0, 0.0, (0.0, 0.0), + vx_dict, _ = modulate_speed(1.0, 0.0, (0.0, 0.0), {"t0": (1.0, 0.0)}) assert math.isclose(vx_list, vx_dict) diff --git a/tests/test_diffdrive.py b/tests/test_diffdrive.py index 27816ac..65844e2 100644 --- a/tests/test_diffdrive.py +++ b/tests/test_diffdrive.py @@ -6,7 +6,7 @@ import pytest from herding.world.diffdrive import ( heading_speed_to_wheels, kinematics_step, - mecanum_inverse, mecanum_kinematics_step, + mecanum_inverse, mecanum_step, velocity_to_mecanum_wheels, velocity_to_wheels, ) @@ -95,7 +95,7 @@ LY = 0.14 # half wheel_base_y def test_mecanum_kinematics_zero_is_identity(): - x, y, h = mecanum_kinematics_step( + x, y, h = mecanum_step( 1.0, 2.0, 0.5, 0.0, 0.0, 0.0, 0.0, WHEEL_R, LX, LY, DT, ) assert (x, y, h) == (1.0, 2.0, 0.5) @@ -104,7 +104,7 @@ def test_mecanum_kinematics_zero_is_identity(): def test_mecanum_kinematics_pure_forward(): # All 4 wheels equal → pure forward (vx_body > 0, vy_body = 0). w = 10.0 - x, y, h = mecanum_kinematics_step( + x, y, h = mecanum_step( 0.0, 0.0, 0.0, w, w, w, w, WHEEL_R, LX, LY, DT, ) assert h == pytest.approx(0.0, abs=1e-9) @@ -118,7 +118,7 @@ def test_mecanum_kinematics_pure_strafe(): # vy_body = (-w_fl+w_fr+w_rl-w_rr)*r/4 > 0 # Use w_fl=-10, w_fr=10, w_rl=10, w_rr=-10. w_fl, w_fr, w_rl, w_rr = -10.0, 10.0, 10.0, -10.0 - x, y, h = mecanum_kinematics_step( + x, y, h = mecanum_step( 0.0, 0.0, 0.0, w_fl, w_fr, w_rl, w_rr, WHEEL_R, LX, LY, DT, ) assert h == pytest.approx(0.0, abs=1e-9) @@ -130,7 +130,7 @@ def test_mecanum_kinematics_pure_strafe(): def test_mecanum_kinematics_strafe_efficiency_scales_y(): # With strafe_efficiency=0.4, realised strafe should be 40% of ideal. w_fl, w_fr, w_rl, w_rr = -10.0, 10.0, 10.0, -10.0 - x, y, _ = mecanum_kinematics_step( + x, y, _ = mecanum_step( 0.0, 0.0, 0.0, w_fl, w_fr, w_rl, w_rr, WHEEL_R, LX, LY, DT, strafe_efficiency=0.4, ) @@ -142,7 +142,7 @@ def test_mecanum_kinematics_strafe_efficiency_scales_y(): def test_mecanum_kinematics_strafe_bleed_pushes_backward(): # Negative bleed means strafe commands also push the body backward. w_fl, w_fr, w_rl, w_rr = -10.0, 10.0, 10.0, -10.0 - x, y, _ = mecanum_kinematics_step( + x, y, _ = mecanum_step( 0.0, 0.0, 0.0, w_fl, w_fr, w_rl, w_rr, WHEEL_R, LX, LY, DT, strafe_efficiency=1.0, strafe_to_forward_bleed=-0.28, @@ -156,7 +156,7 @@ def test_mecanum_kinematics_strafe_bleed_pushes_backward(): def test_mecanum_kinematics_forward_unaffected_by_strafe_params(): # Forward command should be untouched by strafe_efficiency / bleed. w_fl = w_fr = w_rl = w_rr = 10.0 - x, y, _ = mecanum_kinematics_step( + x, y, _ = mecanum_step( 0.0, 0.0, 0.0, w_fl, w_fr, w_rl, w_rr, WHEEL_R, LX, LY, DT, strafe_efficiency=0.4, strafe_to_forward_bleed=-0.28, @@ -170,7 +170,7 @@ def test_mecanum_kinematics_pure_rotation(): # Pure rotation: vx_body=0, vy_body=0, omega>0. # w_fl=-10, w_fr=10, w_rl=-10, w_rr=10 → all sums cancel except omega. w_fl, w_fr, w_rl, w_rr = -10.0, 10.0, -10.0, 10.0 - x, y, h = mecanum_kinematics_step( + x, y, h = mecanum_step( 0.0, 0.0, 0.0, w_fl, w_fr, w_rl, w_rr, WHEEL_R, LX, LY, DT, ) assert x == pytest.approx(0.0, abs=1e-9) diff --git a/tests/test_geometry.py b/tests/test_geometry.py index 7d6a75f..b269752 100644 --- a/tests/test_geometry.py +++ b/tests/test_geometry.py @@ -5,7 +5,7 @@ import math from herding.world.geometry import ( FIELD_X, FIELD_Y, GATE_X, GATE_Y, MAX_SHEEP, PEN_ENTRY, PEN_X, PEN_Y, distance_to_pen_entry, in_field, in_gate_corridor, in_pen, - is_penned_position, + is_penned, ) @@ -44,23 +44,23 @@ def test_in_gate_corridor(): assert not in_gate_corridor(5.0, -18.0) -def test_is_penned_position_latches_below_gate(): +def test_is_penned_latches_below_gate(): # In the gate column and south of the gate plane → penned. - assert is_penned_position(11.5, -15.0) - assert is_penned_position(10.5, -18.0) - assert is_penned_position(12.5, -22.0) + assert is_penned(11.5, -15.0) + assert is_penned(10.5, -18.0) + assert is_penned(12.5, -22.0) # Above the gate plane → not yet. - assert not is_penned_position(11.5, -14.9) + assert not is_penned(11.5, -14.9) # Outside the gate column → not penned even if south. - assert not is_penned_position(0.0, -16.0) - assert not is_penned_position(14.0, -16.0) + assert not is_penned(0.0, -16.0) + assert not is_penned(14.0, -16.0) -def test_is_penned_position_latch_margin(): +def test_is_penned_latch_margin(): # Slight tolerance on the gate column. - assert is_penned_position(9.9, -15.5) - assert is_penned_position(13.1, -15.5) - assert not is_penned_position(9.7, -15.5) + assert is_penned(9.9, -15.5) + assert is_penned(13.1, -15.5) + assert not is_penned(9.7, -15.5) def test_distance_to_pen_entry(): diff --git a/tests/test_perception.py b/tests/test_perception.py index fc137c7..229ade2 100644 --- a/tests/test_perception.py +++ b/tests/test_perception.py @@ -136,7 +136,7 @@ def test_tracker_forgets_stale_tracks(): def test_tracker_penned_position_promotes_track(): t = SheepTracker() t.update([(11.5, -16.0)]) # spawn inside the pen column - # is_penned_position is True for this point. + # is_penned is True for this point. assert t.n_penned() == 1 assert t.n_active() == 0 diff --git a/training/bc/pretrain.py b/training/bc/pretrain.py index ecb3303..870729e 100644 --- a/training/bc/pretrain.py +++ b/training/bc/pretrain.py @@ -55,7 +55,7 @@ def build_model(net_arch_pi, net_arch_vf, log_std_init: float, return model, env -def policy_forward_mean(policy, obs_batch): +def forward_mean(policy, obs_batch): """Return the deterministic mean action for an obs batch. SB3's ActorCriticPolicy routes ``forward`` through a Distribution @@ -177,7 +177,7 @@ def main(): ob_batch = ob_batch.to(args.device) act_batch = act_batch.to(args.device) optimizer.zero_grad() - mean_action = policy_forward_mean(policy, ob_batch) + mean_action = forward_mean(policy, ob_batch) loss, mse_val, cos_val = combined_loss(mean_action, act_batch) loss.backward() optimizer.step() @@ -196,7 +196,7 @@ def main(): for ob_batch, act_batch in val_loader: ob_batch = ob_batch.to(args.device) act_batch = act_batch.to(args.device) - mean_action = policy_forward_mean(policy, ob_batch) + mean_action = forward_mean(policy, ob_batch) bs = ob_batch.size(0) val_total += nn.functional.mse_loss( mean_action, act_batch, reduction="sum", diff --git a/training/herding_env.py b/training/herding_env.py index 7719e5c..d0d20a2 100644 --- a/training/herding_env.py +++ b/training/herding_env.py @@ -28,7 +28,7 @@ from gymnasium import spaces from herding.world.diffdrive import ( heading_speed_to_wheels, kinematics_step, - mecanum_kinematics_step, velocity_to_mecanum_wheels, velocity_to_wheels, + mecanum_step, velocity_to_mecanum_wheels, velocity_to_wheels, ) from herding.world.flocking_sim import ( FLEE_SPEED, MAX_SPEED, WANDER_SPEED, compute_heading_speed, @@ -40,7 +40,7 @@ from herding.world.geometry import ( GATE_X, GATE_Y, MAX_SHEEP, PEN_ENTRY, PEN_X, PEN_Y, SHEEP_MAX_WHEEL_OMEGA, SHEEP_WHEEL_BASE, SHEEP_WHEEL_RADIUS, - WEBOTS_DT, clip_to_field, is_penned_position, + WEBOTS_DT, clip_to_field, is_penned, ) from herding.perception.lidar_perception import detections_from_scan from herding.perception.lidar_sim import simulate_scan @@ -302,7 +302,7 @@ class HerdingEnv(gym.Env): if robot_cfg is not None else 1.0) strafe_bleed = (robot_cfg.strafe_to_forward_bleed if robot_cfg is not None else 0.0) - self.dog_x, self.dog_y, self.dog_heading = mecanum_kinematics_step( + self.dog_x, self.dog_y, self.dog_heading = mecanum_step( self.dog_x, self.dog_y, self.dog_heading, w_fl, w_fr, w_rl, w_rr, DOG_WHEEL_RADIUS, @@ -337,7 +337,7 @@ class HerdingEnv(gym.Env): self._step_one_sheep(i) for i in range(self.n_sheep): if (not self.sheep_penned[i] - and is_penned_position(self.sheep_x[i], self.sheep_y[i])): + and is_penned(self.sheep_x[i], self.sheep_y[i])): self.sheep_penned[i] = True # LiDAR perception runs after sheep move; feeds the obs and the