♻️ crazy refactor
This commit is contained in:
381
src/sysid/capture.py
Normal file
381
src/sysid/capture.py
Normal file
@@ -0,0 +1,381 @@
|
||||
"""Capture a real-robot trajectory under random excitation (PRBS-style).
|
||||
|
||||
Connects to the ESP32 over serial, sends random PWM commands to excite
|
||||
the system, and records motor + pendulum angles and velocities at ~50 Hz.
|
||||
|
||||
Saves a compressed numpy archive (.npz) that the optimizer can replay
|
||||
in simulation to fit physics parameters.
|
||||
|
||||
Usage:
|
||||
python -m src.sysid.capture \
|
||||
--robot-path assets/rotary_cartpole \
|
||||
--port /dev/cu.usbserial-0001 \
|
||||
--duration 20
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import math
|
||||
import os
|
||||
import random
|
||||
import threading
|
||||
import time
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import structlog
|
||||
import yaml
|
||||
|
||||
log = structlog.get_logger()
|
||||
|
||||
|
||||
# ── Serial protocol helpers (mirrored from SerialRunner) ─────────────
|
||||
|
||||
|
||||
def _parse_state_line(line: str) -> dict[str, Any] | None:
|
||||
"""Parse an ``S,…`` state line from the ESP32."""
|
||||
if not line.startswith("S,"):
|
||||
return None
|
||||
parts = line.split(",")
|
||||
if len(parts) < 12:
|
||||
return None
|
||||
try:
|
||||
return {
|
||||
"timestamp_ms": int(parts[1]),
|
||||
"encoder_count": int(parts[2]),
|
||||
"rpm": float(parts[3]),
|
||||
"motor_speed": int(parts[4]),
|
||||
"at_limit": bool(int(parts[5])),
|
||||
"pendulum_angle": float(parts[6]),
|
||||
"pendulum_velocity": float(parts[7]),
|
||||
"target_speed": int(parts[8]),
|
||||
"braking": bool(int(parts[9])),
|
||||
"enc_vel_cps": float(parts[10]),
|
||||
"pendulum_ok": bool(int(parts[11])),
|
||||
}
|
||||
except (ValueError, IndexError):
|
||||
return None
|
||||
|
||||
|
||||
# ── Background serial reader ─────────────────────────────────────────
|
||||
|
||||
|
||||
class _SerialReader:
|
||||
"""Minimal background reader for the ESP32 serial stream."""
|
||||
|
||||
def __init__(self, port: str, baud: int = 115200):
|
||||
import serial as _serial
|
||||
|
||||
self._serial_mod = _serial
|
||||
self.ser = _serial.Serial(port, baud, timeout=0.05)
|
||||
time.sleep(2) # Wait for ESP32 boot.
|
||||
self.ser.reset_input_buffer()
|
||||
|
||||
self._latest: dict[str, Any] = {}
|
||||
self._lock = threading.Lock()
|
||||
self._event = threading.Event()
|
||||
self._running = True
|
||||
|
||||
self._thread = threading.Thread(target=self._reader_loop, daemon=True)
|
||||
self._thread.start()
|
||||
|
||||
def _reader_loop(self) -> None:
|
||||
while self._running:
|
||||
try:
|
||||
if self.ser.in_waiting:
|
||||
line = (
|
||||
self.ser.readline()
|
||||
.decode("utf-8", errors="ignore")
|
||||
.strip()
|
||||
)
|
||||
parsed = _parse_state_line(line)
|
||||
if parsed is not None:
|
||||
with self._lock:
|
||||
self._latest = parsed
|
||||
self._event.set()
|
||||
else:
|
||||
time.sleep(0.001)
|
||||
except (OSError, self._serial_mod.SerialException):
|
||||
log.critical("serial_lost")
|
||||
break
|
||||
|
||||
def send(self, cmd: str) -> None:
|
||||
try:
|
||||
self.ser.write(f"{cmd}\n".encode())
|
||||
except (OSError, self._serial_mod.SerialException):
|
||||
log.critical("serial_send_failed", cmd=cmd)
|
||||
|
||||
def read(self) -> dict[str, Any]:
|
||||
with self._lock:
|
||||
return dict(self._latest)
|
||||
|
||||
def read_blocking(self, timeout: float = 0.1) -> dict[str, Any]:
|
||||
self._event.clear()
|
||||
self._event.wait(timeout=timeout)
|
||||
return self.read()
|
||||
|
||||
def close(self) -> None:
|
||||
self._running = False
|
||||
self.send("H")
|
||||
self.send("M0")
|
||||
time.sleep(0.1)
|
||||
self._thread.join(timeout=1.0)
|
||||
self.ser.close()
|
||||
|
||||
|
||||
# ── PRBS excitation signal ───────────────────────────────────────────
|
||||
|
||||
|
||||
class _PRBSExcitation:
|
||||
"""Random hold-value excitation with configurable amplitude and hold time.
|
||||
|
||||
At each call to ``__call__``, returns the current PWM value.
|
||||
The value is held for a random duration (``hold_min``–``hold_max`` ms),
|
||||
then a new random value is drawn uniformly from ``[-amplitude, +amplitude]``.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
amplitude: int = 180,
|
||||
hold_min_ms: int = 50,
|
||||
hold_max_ms: int = 300,
|
||||
):
|
||||
self.amplitude = amplitude
|
||||
self.hold_min_ms = hold_min_ms
|
||||
self.hold_max_ms = hold_max_ms
|
||||
self._current: int = 0
|
||||
self._switch_time: float = 0.0
|
||||
self._new_value()
|
||||
|
||||
def _new_value(self) -> None:
|
||||
self._current = random.randint(-self.amplitude, self.amplitude)
|
||||
hold_ms = random.randint(self.hold_min_ms, self.hold_max_ms)
|
||||
self._switch_time = time.monotonic() + hold_ms / 1000.0
|
||||
|
||||
def __call__(self) -> int:
|
||||
if time.monotonic() >= self._switch_time:
|
||||
self._new_value()
|
||||
return self._current
|
||||
|
||||
|
||||
# ── Main capture loop ────────────────────────────────────────────────
|
||||
|
||||
|
||||
def capture(
|
||||
robot_path: str | Path,
|
||||
port: str = "/dev/cu.usbserial-0001",
|
||||
baud: int = 115200,
|
||||
duration: float = 20.0,
|
||||
amplitude: int = 180,
|
||||
hold_min_ms: int = 50,
|
||||
hold_max_ms: int = 300,
|
||||
dt: float = 0.02,
|
||||
) -> Path:
|
||||
"""Run the capture procedure and return the path to the saved .npz file.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
robot_path : path to robot asset directory (contains hardware.yaml)
|
||||
port : serial port for ESP32
|
||||
baud : baud rate
|
||||
duration : capture duration in seconds
|
||||
amplitude : max PWM magnitude for excitation (0–255)
|
||||
hold_min_ms / hold_max_ms : random hold time range (ms)
|
||||
dt : target sampling period (seconds), default 50 Hz
|
||||
"""
|
||||
robot_path = Path(robot_path).resolve()
|
||||
|
||||
# Load hardware config for encoder conversion + safety.
|
||||
hw_yaml = robot_path / "hardware.yaml"
|
||||
if not hw_yaml.exists():
|
||||
raise FileNotFoundError(f"hardware.yaml not found in {robot_path}")
|
||||
raw_hw = yaml.safe_load(hw_yaml.read_text())
|
||||
ppr = raw_hw.get("encoder", {}).get("ppr", 11)
|
||||
gear_ratio = raw_hw.get("encoder", {}).get("gear_ratio", 30.0)
|
||||
counts_per_rev: float = ppr * gear_ratio * 4.0
|
||||
max_motor_deg = raw_hw.get("safety", {}).get("max_motor_angle_deg", 90.0)
|
||||
max_motor_rad = math.radians(max_motor_deg) if max_motor_deg > 0 else 0.0
|
||||
|
||||
# Connect.
|
||||
reader = _SerialReader(port, baud)
|
||||
excitation = _PRBSExcitation(amplitude, hold_min_ms, hold_max_ms)
|
||||
|
||||
# Prepare recording buffers.
|
||||
max_samples = int(duration / dt) + 500 # headroom
|
||||
rec_time = np.zeros(max_samples, dtype=np.float64)
|
||||
rec_action = np.zeros(max_samples, dtype=np.float64)
|
||||
rec_motor_angle = np.zeros(max_samples, dtype=np.float64)
|
||||
rec_motor_vel = np.zeros(max_samples, dtype=np.float64)
|
||||
rec_pend_angle = np.zeros(max_samples, dtype=np.float64)
|
||||
rec_pend_vel = np.zeros(max_samples, dtype=np.float64)
|
||||
|
||||
# Start streaming.
|
||||
reader.send("G")
|
||||
time.sleep(0.1)
|
||||
|
||||
log.info(
|
||||
"capture_starting",
|
||||
port=port,
|
||||
duration=duration,
|
||||
amplitude=amplitude,
|
||||
hold_range_ms=f"{hold_min_ms}–{hold_max_ms}",
|
||||
dt=dt,
|
||||
)
|
||||
|
||||
idx = 0
|
||||
t0 = time.monotonic()
|
||||
try:
|
||||
while True:
|
||||
loop_start = time.monotonic()
|
||||
elapsed = loop_start - t0
|
||||
if elapsed >= duration:
|
||||
break
|
||||
|
||||
# Get excitation PWM.
|
||||
pwm = excitation()
|
||||
|
||||
# Safety: reverse/zero if near motor limit.
|
||||
state = reader.read()
|
||||
if state:
|
||||
motor_angle_rad = (
|
||||
state.get("encoder_count", 0) / counts_per_rev * 2.0 * math.pi
|
||||
)
|
||||
if max_motor_rad > 0:
|
||||
margin = max_motor_rad * 0.85 # start braking at 85%
|
||||
if motor_angle_rad > margin and pwm > 0:
|
||||
pwm = -abs(pwm) # reverse
|
||||
elif motor_angle_rad < -margin and pwm < 0:
|
||||
pwm = abs(pwm) # reverse
|
||||
|
||||
# Send command.
|
||||
reader.send(f"M{pwm}")
|
||||
|
||||
# Wait for fresh data.
|
||||
time.sleep(max(0, dt - (time.monotonic() - loop_start) - 0.005))
|
||||
state = reader.read_blocking(timeout=dt)
|
||||
|
||||
if state:
|
||||
enc = state.get("encoder_count", 0)
|
||||
motor_angle = enc / counts_per_rev * 2.0 * math.pi
|
||||
motor_vel = (
|
||||
state.get("enc_vel_cps", 0.0) / counts_per_rev * 2.0 * math.pi
|
||||
)
|
||||
pend_angle = math.radians(state.get("pendulum_angle", 0.0))
|
||||
pend_vel = math.radians(state.get("pendulum_velocity", 0.0))
|
||||
|
||||
# Normalised action: PWM / 255 → [-1, 1]
|
||||
action_norm = pwm / 255.0
|
||||
|
||||
if idx < max_samples:
|
||||
rec_time[idx] = elapsed
|
||||
rec_action[idx] = action_norm
|
||||
rec_motor_angle[idx] = motor_angle
|
||||
rec_motor_vel[idx] = motor_vel
|
||||
rec_pend_angle[idx] = pend_angle
|
||||
rec_pend_vel[idx] = pend_vel
|
||||
idx += 1
|
||||
|
||||
# Progress.
|
||||
if idx % 50 == 0:
|
||||
log.info(
|
||||
"capture_progress",
|
||||
elapsed=f"{elapsed:.1f}/{duration:.0f}s",
|
||||
samples=idx,
|
||||
pwm=pwm,
|
||||
)
|
||||
|
||||
# Pace to dt.
|
||||
remaining = dt - (time.monotonic() - loop_start)
|
||||
if remaining > 0:
|
||||
time.sleep(remaining)
|
||||
|
||||
finally:
|
||||
reader.send("M0")
|
||||
reader.close()
|
||||
|
||||
# Trim to actual sample count.
|
||||
rec_time = rec_time[:idx]
|
||||
rec_action = rec_action[:idx]
|
||||
rec_motor_angle = rec_motor_angle[:idx]
|
||||
rec_motor_vel = rec_motor_vel[:idx]
|
||||
rec_pend_angle = rec_pend_angle[:idx]
|
||||
rec_pend_vel = rec_pend_vel[:idx]
|
||||
|
||||
# Save.
|
||||
recordings_dir = robot_path / "recordings"
|
||||
recordings_dir.mkdir(exist_ok=True)
|
||||
stamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
out_path = recordings_dir / f"capture_{stamp}.npz"
|
||||
np.savez_compressed(
|
||||
out_path,
|
||||
time=rec_time,
|
||||
action=rec_action,
|
||||
motor_angle=rec_motor_angle,
|
||||
motor_vel=rec_motor_vel,
|
||||
pendulum_angle=rec_pend_angle,
|
||||
pendulum_vel=rec_pend_vel,
|
||||
)
|
||||
|
||||
log.info(
|
||||
"capture_saved",
|
||||
path=str(out_path),
|
||||
samples=idx,
|
||||
duration_actual=f"{rec_time[-1]:.2f}s" if idx > 0 else "0s",
|
||||
)
|
||||
return out_path
|
||||
|
||||
|
||||
# ── CLI entry point ──────────────────────────────────────────────────
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Capture a real-robot trajectory for system identification."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--robot-path",
|
||||
type=str,
|
||||
default="assets/rotary_cartpole",
|
||||
help="Path to robot asset directory (contains hardware.yaml)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--port",
|
||||
type=str,
|
||||
default="/dev/cu.usbserial-0001",
|
||||
help="Serial port for ESP32",
|
||||
)
|
||||
parser.add_argument("--baud", type=int, default=115200)
|
||||
parser.add_argument(
|
||||
"--duration", type=float, default=20.0, help="Capture duration (s)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--amplitude", type=int, default=180, help="Max PWM magnitude (0–255)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--hold-min-ms", type=int, default=50, help="Min hold time (ms)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--hold-max-ms", type=int, default=300, help="Max hold time (ms)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dt", type=float, default=0.02, help="Sample period (s)"
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
capture(
|
||||
robot_path=args.robot_path,
|
||||
port=args.port,
|
||||
baud=args.baud,
|
||||
duration=args.duration,
|
||||
amplitude=args.amplitude,
|
||||
hold_min_ms=args.hold_min_ms,
|
||||
hold_max_ms=args.hold_max_ms,
|
||||
dt=args.dt,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user