1"""Causal sampled control track for inverse-gain structured distillation."""
 2import numpy as np
 3
 4META = {
 5    "name": "inverse_gain_control",
 6    "domain": "dynamics",
 7    "description": "Causal windows of regulated state, velocity, and reference; target is privileged inverse-gain expert action."
 8}
 9
10def get_dataset(seed, n_train=400, n_test=400):
11    def make(n, offset):
12        rng = np.random.RandomState(seed + offset)
13        dt, k1 = 0.05, 1.5
14        X = np.empty((n, 24), np.float32); Y = np.empty((n, 1), np.float32)
15        for i in range(n):
16            a = rng.choice([0.25, 0.5, 1.0, 2.0])
17            d = rng.normal(0, 0.35)
18            x1, x2 = rng.normal(0, .3), rng.normal(0, .2)
19            hist = []
20            for k in range(8):
21                r = 1.0*np.sin(.045*k + rng.uniform(0, 6.28)) + .25*np.sin(.13*k)
22                prev = x1
23                # sampled plant: x1 evolves under unknown gain and disturbance
24                x1 = x1 + dt*(a*x2 + d)
25                delta = (x1-prev)/dt
26                e = r-x1
27                expert = x2 + (k1*e-delta)/a
28                hist.append([x1, x2, r])
29                # low-level actuator tracks the expert reference
30                x2 = x2 + dt*3.0*(expert-x2)
31            X[i] = np.asarray(hist, np.float32).reshape(-1)
32            Y[i, 0] = expert
33        return X, Y
34    xtr, ytr = make(n_train, 0); xte, yte = make(n_test, 5000)
35    return {"xtr": xtr, "ytr": ytr, "xte": xte, "yte": yte,
36            "task": "regression", "metric": "mse"}