dronepilot: first upload

8d8c164 verified 8 days ago

2.28 kB

"""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.")

	"""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.")