1import numpy as np
 2
 3META = {
 4    "name": "multi_agent_critic_compatibility",
 5    "domain": "dynamics/control",
 6    "description": "Two-agent centralized critic regression on joint continuous actions with cross-partial compatibility structure."
 7}
 8
 9def get_dataset(seed, n_train, n_test):
10    rng = np.random.RandomState(seed)
11    def sample(n, noisy):
12        x = rng.uniform(-1.0, 1.0, size=(n, 6)).astype(np.float32)
13        s0, s1, a, b = x[:, 0], x[:, 1], x[:, 2], x[:, 4]
14        common = 0.45*a*b + 0.18*np.sin(a + 0.3*s0)*np.sin(b - 0.2*s1)
15        q1 = 0.55*s0*a - 0.25*s1*b + common + 0.20*a*a - 0.12*b*b
16        q2 = -0.30*s0*a + 0.48*s1*b + common - 0.10*a*a + 0.16*b*b
17        if noisy:
18            q1 += rng.normal(0, 0.16, n) + 0.22*a*b
19            q2 += rng.normal(0, 0.16, n) - 0.22*a*b
20        return x, np.stack([q1, q2], axis=1).astype(np.float32)
21    xtr, ytr = sample(n_train, True)
22    xte, yte = sample(n_test, False)
23    return {"xtr": xtr, "ytr": ytr, "xte": xte, "yte": yte,
24            "task": "regression", "metric": "mse", "out_dim": 2}