File size: 2,280 Bytes
8d8c164 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 | """Imagine a plane's next 2 seconds under different commands, with this world model and with physics alone.
pip install torch numpy pyyaml
python example.py
"""
from pathlib import Path
import numpy as np
import torch
import yaml
np.set_printoptions(suppress=True, precision=2)
from fixed_wing.dynamics import Airframe
from fixed_wing.learned import LearnedFixedWing, rollout_with
HERE = Path(__file__).parent
cfg = yaml.safe_load((HERE / "fixed_wing/airframe.yaml").read_text())
wm = LearnedFixedWing.load(HERE / "fw_wm_small.pt", cfg)
print(f"world model: {sum(p.numel() for p in wm.corr.parameters())} learned parameters on a 6-DoF physics backbone")
# shared state layout (18): position (ENU, m), velocity (m/s), rotation matrix columns (body FLU -> world), body rates
TRIM = np.array([0.613, 0.0, -0.0714, 0.0]) # throttle, aileron, elevator, rudder: level flight at 14 m/s
alpha = 0.031 # trim angle of attack (rad): the nose sits slightly up
R = np.array([[np.cos(alpha), 0, -np.sin(alpha)], [0, 1, 0], [np.sin(alpha), 0, np.cos(alpha)]])
x = np.concatenate([[0, 0, 40.0], [14.0, 0, 0], R[:, 0], R[:, 1], R[:, 2], [0, 0, 0]])
history = np.repeat(x[None], wm.H + 1, 0) # 0.5 s of straight, level flight
past = np.repeat(TRIM[None], wm.H, 0)
T = 40 # 2 s at 20 Hz
plans = {"hold trim": np.repeat(TRIM[None], T, 0),
"bank right": np.repeat((TRIM + [0, 0.10, 0, 0])[None], T, 0),
"pull up": np.repeat((TRIM + [0.2, 0, -0.08, 0])[None], T, 0)}
futures = wm.rollout(history, past, np.stack(list(plans.values()))) # (3, 40, 18)
physics = rollout_with(Airframe.from_cfg(cfg), None, None, torch.tensor(history[-1:], dtype=torch.float32).expand(3, -1),
torch.tensor(np.stack(list(plans.values())), dtype=torch.float32))
for k, name in enumerate(plans):
p, q = futures[k, -1, :3].numpy(), physics[k, -1, :3].numpy()
print(f"{name:10s} world model -> {np.round(p, 2)} physics alone -> {np.round(q, 2)}")
print("On an airframe exactly like the nominal one in still air, the two agree closely; the world model earns its keep "
"when the history shows a different airframe, wind, or lag.")
|