import numpy as np

META = {
    "name": "multi_agent_critic_compatibility",
    "domain": "dynamics/control",
    "description": "Two-agent centralized critic regression on joint continuous actions with cross-partial compatibility structure."
}

def get_dataset(seed, n_train, n_test):
    rng = np.random.RandomState(seed)
    def sample(n, noisy):
        x = rng.uniform(-1.0, 1.0, size=(n, 6)).astype(np.float32)
        s0, s1, a, b = x[:, 0], x[:, 1], x[:, 2], x[:, 4]
        common = 0.45*a*b + 0.18*np.sin(a + 0.3*s0)*np.sin(b - 0.2*s1)
        q1 = 0.55*s0*a - 0.25*s1*b + common + 0.20*a*a - 0.12*b*b
        q2 = -0.30*s0*a + 0.48*s1*b + common - 0.10*a*a + 0.16*b*b
        if noisy:
            q1 += rng.normal(0, 0.16, n) + 0.22*a*b
            q2 += rng.normal(0, 0.16, n) - 0.22*a*b
        return x, np.stack([q1, q2], axis=1).astype(np.float32)
    xtr, ytr = sample(n_train, True)
    xte, yte = sample(n_test, False)
    return {"xtr": xtr, "ytr": ytr, "xte": xte, "yte": yte,
            "task": "regression", "metric": "mse", "out_dim": 2}