Frank14fandCursor
61e82ec90a
Add deterministic phase-filtered V5 and Legacy acquisition with complete-cycle mean fields, then evaluate controlled wakes against target and zero baselines offline. Co-authored-by: Cursor <cursoragent@cursor.com>
233 lines
9.5 KiB
Python
233 lines
9.5 KiB
Python
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
import src.drl_pinball.legacy_test.acquire as acquire
|
|
import src.drl_pinball.legacy_test.core.legacy_env_builder as builder
|
|
|
|
|
|
class RoleFlow:
|
|
DATA_TYPE = np.float32
|
|
|
|
def __init__(self, n_objects, obs_size, absolute_base=2_000_000):
|
|
self.action = np.ones(n_objects, dtype=np.float32)
|
|
self.obs = np.zeros(obs_size, dtype=np.float32)
|
|
self.intervals = 0
|
|
self.absolute_base = absolute_base
|
|
self.commands = []
|
|
self.capture_intervals = []
|
|
|
|
def restore_ddf(self):
|
|
pass
|
|
|
|
def apply_ddf(self):
|
|
pass
|
|
|
|
def run(self, steps, command):
|
|
assert steps == acquire.SI
|
|
assert np.array_equal(self.action, np.zeros_like(self.action))
|
|
self.intervals += 1
|
|
self.commands.append(command.copy())
|
|
self.action = command.copy()
|
|
phase = 2 * np.pi * self.intervals / acquire.CYCLE_LENGTH
|
|
self.obs[0:2] = [1000 + self.intervals, 2000 + self.intervals]
|
|
self.obs[2:8] = [np.sin(phase + index) for index in range(6)]
|
|
if self.obs.size >= 14:
|
|
self.obs[8:14] = np.arange(6) + self.intervals
|
|
|
|
def current_effective_action(self):
|
|
return self.action.copy()
|
|
|
|
def solver_clock_state(self):
|
|
return {"solver_absolute_lattice_clock": self.absolute_base + self.intervals * acquire.SI,
|
|
"solver_absolute_control_clock": self.intervals}
|
|
|
|
|
|
def data(n_objects=7):
|
|
phase = 2 * np.pi * np.arange(150) / acquire.CYCLE_LENGTH
|
|
target = np.column_stack([np.sin(phase + index) for index in range(6)]).astype(np.float32)
|
|
result = {
|
|
"config": {"nx": 1280, "ny": 2, "sample_interval": 800,
|
|
"n_obj_total": n_objects},
|
|
"target_states": target,
|
|
}
|
|
if n_objects == 7:
|
|
result["norm"] = {"save_states": np.zeros((150, 12), dtype=np.float32)}
|
|
return result
|
|
|
|
|
|
def norm():
|
|
return {"force_norm_fact": 1.0, "sens_deviation": np.zeros(6),
|
|
"sens_norm_fact": np.ones(6)}
|
|
|
|
|
|
def capture(flow, nx, ny):
|
|
assert flow.intervals > acquire.WARMUP_INTERVALS
|
|
flow.capture_intervals.append(flow.intervals)
|
|
return {"ux": np.zeros((ny, nx), np.float32), "uy": np.zeros((ny, nx), np.float32)}
|
|
|
|
|
|
def test_target_builder_uses_disturbance_then_three_sensors_without_pinball(monkeypatch):
|
|
class TargetFlow:
|
|
DATA_TYPE = np.float32
|
|
FIELD_SHAPE = (1280, 2, 1)
|
|
|
|
def __init__(self, field_cfg, cuda_cfg, device_id):
|
|
self.obs = np.zeros(0, np.float32)
|
|
self.objects = []
|
|
self.interval = 0
|
|
|
|
def add_cylinder(self, center, radius):
|
|
self.objects.append(("cylinder", center, radius))
|
|
self.obs = np.zeros(2 * len(self.objects), np.float32)
|
|
|
|
def add_sensor(self, center, radius):
|
|
self.objects.append(("sensor", center, radius))
|
|
self.obs = np.zeros(2 * len(self.objects), np.float32)
|
|
|
|
class Config:
|
|
def _replace(self, **kwargs):
|
|
return self
|
|
|
|
monkeypatch.setattr(builder, "FlowField", TargetFlow)
|
|
monkeypatch.setattr(builder.legacy_utils, "load_cuda_config", lambda path: object())
|
|
monkeypatch.setattr(builder.legacy_utils, "load_flow_field_config", lambda path: Config())
|
|
monkeypatch.setattr(builder, "_stabilize", lambda flow, count: None)
|
|
monkeypatch.setattr(builder, "FIFO_LEN", 2)
|
|
|
|
def interval(flow, steps, command):
|
|
flow.interval += 1
|
|
flow.obs[0:2] = [100 + flow.interval, 200 + flow.interval]
|
|
flow.obs[2:8] = np.arange(6) + 10 * flow.interval
|
|
|
|
monkeypatch.setattr(builder, "run_historical_interval", interval)
|
|
result = builder.build_karman_target(0, 100.0, sample_interval=800)
|
|
|
|
flow = result["flow_field"]
|
|
assert [item[0] for item in flow.objects] == ["cylinder", "sensor", "sensor", "sensor"]
|
|
assert flow.objects[0][1][0] == 10.0 * builder.L0
|
|
assert [item[1][1] for item in flow.objects[1:]] == [40.5, 0.5, -39.5]
|
|
assert result["config"]["n_obj_total"] == 4
|
|
assert np.array_equal(result["target_states"],
|
|
np.asarray([np.arange(6) + 10, np.arange(6) + 20], np.float32))
|
|
|
|
|
|
def test_zero_role_schedule_actions_rewards_and_fields(monkeypatch, tmp_path):
|
|
flow = RoleFlow(7, 14)
|
|
reward_calls = []
|
|
|
|
def rewards(*args):
|
|
reward_calls.append(args)
|
|
return {"reward": 0.4, "reward_cd": 0.5, "reward_cl": 0.6,
|
|
"native_legacy_dtw": 0.7}
|
|
|
|
monkeypatch.setattr(acquire, "reward_terms", rewards)
|
|
rows, fields = acquire._collect_role(
|
|
"zero", data(), None, flow, norm(), tmp_path, capture_field=capture,
|
|
)
|
|
|
|
assert flow.intervals == len(reward_calls) == 640
|
|
assert flow.capture_intervals == list(range(481, 641))
|
|
assert len(rows) == 160
|
|
assert fields["ux"].shape == fields["uy"].shape == (160, 2, 1280)
|
|
assert not list(tmp_path.glob("boundary_*.npz"))
|
|
counter_bias = np.array([0.0, 0.5, -0.5], np.float32)
|
|
assert np.array_equal(rows[0]["action_normalized"], counter_bias)
|
|
assert all(np.array_equal(command, np.zeros(7, np.float32)) for command in flow.commands)
|
|
assert np.array_equal(rows[0]["commanded_target_omega"], np.zeros(3, np.float32))
|
|
assert np.allclose(rows[0]["effective_smoothed_omega"], 0.0, atol=1e-7)
|
|
assert rows[0]["reward_raw"] == 0.4
|
|
assert rows[0]["native_reward_dtw"] == 0.7
|
|
assert rows[0]["control_index"] == 481 and rows[-1]["control_index"] == 640
|
|
|
|
|
|
def test_zero_role_metadata_names_physical_zero_counter_bias():
|
|
source = Path(acquire.__file__).read_text(encoding="utf-8")
|
|
assert "physical-zero/uncontrolled" in source
|
|
assert "counter-bias normalized action" in source
|
|
|
|
|
|
def test_target_role_ordering_nan_contract_and_schedule(monkeypatch, tmp_path):
|
|
flow = RoleFlow(4, 8)
|
|
monkeypatch.setattr(acquire, "reward_terms", lambda *args: pytest.fail("target has no reward"))
|
|
rows, fields = acquire._collect_role(
|
|
"target", data(n_objects=4), None, flow, None, tmp_path, capture_field=capture,
|
|
)
|
|
|
|
assert flow.intervals == 640
|
|
assert all(np.array_equal(command, np.zeros(4, np.float32)) for command in flow.commands)
|
|
assert len(rows) == 160
|
|
assert fields["ux"].shape == (160, 2, 1280)
|
|
first = rows[0]
|
|
phase = 2 * np.pi * 481 / acquire.CYCLE_LENGTH
|
|
expected = np.asarray([np.sin(phase + index) for index in range(6)], np.float32)
|
|
assert np.allclose(first["sensors"], expected)
|
|
assert first["sensors"][0] != flow.obs[0]
|
|
for name in ("forces", "action_normalized", "commanded_target_omega",
|
|
"effective_smoothed_omega"):
|
|
assert np.isnan(first[name]).all()
|
|
for name in ("reward_raw", "reward_cd", "reward_cl", "reward_sim",
|
|
"native_reward_dtw"):
|
|
assert np.isnan(first[name])
|
|
assert first["solver_absolute_lattice_step"] == 2_000_000 + 481 * 800
|
|
|
|
|
|
def test_target_sensor_reference_selects_legacy_target_channels():
|
|
six = np.arange(60, dtype=float).reshape(10, 6)
|
|
eight = np.column_stack((np.full((10, 2), -1.0), six))
|
|
assert np.array_equal(acquire._target_sensors(six), six)
|
|
assert np.array_equal(acquire._target_sensors(eight), six)
|
|
with pytest.raises(ValueError, match="six sensor channels"):
|
|
acquire._target_sensors(np.zeros((10, 7)))
|
|
|
|
|
|
def test_target_phase_artifacts_exclude_unavailable_nan_variables():
|
|
columns = {"sensors": np.ones((4, 6)), "forces": np.full((4, 6), np.nan),
|
|
"action_normalized": np.full((4, 3), np.nan),
|
|
"physical_time": np.arange(4), "lattice_step": np.arange(4),
|
|
"solver_absolute_lattice_step": np.arange(4),
|
|
"control_index": np.arange(4), "phase": np.arange(4),
|
|
"cycle_id": np.arange(4)}
|
|
values = acquire._phase_values(columns, "target")
|
|
assert set(values) == {"sensors"}
|
|
assert np.isfinite(values["sensors"]).all()
|
|
|
|
|
|
@pytest.mark.parametrize("role", ["target", "zero"])
|
|
def test_role_failure_transaction_preserves_existing(monkeypatch, tmp_path, role):
|
|
root = tmp_path / "out"
|
|
final = root / "legacy" / acquire.CASE_NAME / role
|
|
final.mkdir(parents=True)
|
|
marker = final / "keep.txt"
|
|
marker.write_text("old", encoding="utf-8")
|
|
model = tmp_path / "model.zip"; model.write_bytes(b"model")
|
|
norm_path = tmp_path / "norm.json"; norm_path.write_text("{}", encoding="utf-8")
|
|
bundle = {"case": acquire.get_case(acquire.CASE_NAME), "model_path": model,
|
|
"norm_path": norm_path, "norm": norm(),
|
|
"reference_path": tmp_path / "missing.npz"}
|
|
monkeypatch.setattr(acquire, "_resolve_bundle", lambda repo, *args: bundle)
|
|
monkeypatch.setattr(acquire, "_collect_role", lambda *args, **kwargs: ([], []))
|
|
|
|
def storage_validator(**kwargs):
|
|
root.mkdir(exist_ok=True)
|
|
return {"resolved_output_root": root, "device": root.stat().st_dev}
|
|
|
|
def runtime_factory(case, current_bundle, device, current_role):
|
|
n_objects = 4 if current_role == "target" else 7
|
|
return data(n_objects), None, RoleFlow(n_objects, 8 if current_role == "target" else 14)
|
|
|
|
with pytest.raises(RuntimeError, match="finalize failed"):
|
|
acquire.acquire_role(
|
|
role, output_root=root, overwrite=True, repo_root=tmp_path,
|
|
storage_validator=storage_validator, runtime_factory=runtime_factory,
|
|
finalizer=lambda *args: (_ for _ in ()).throw(RuntimeError("finalize failed")),
|
|
)
|
|
assert marker.read_text(encoding="utf-8") == "old"
|
|
assert not list(final.parent.glob(f".{role}.staging-*"))
|
|
|
|
|
|
def test_cli_exposes_all_roles():
|
|
source = Path(acquire.__file__).read_text(encoding="utf-8")
|
|
assert 'choices=ROLES' in source
|