1import numpy as np
 2META = {'name':'joint_action_values','domain':'multi_agent_rl','description':'Two-player continuous action values with asymmetric cross-effects and cooperative state-dependent optima.'}
 3def get_dataset(seed, n_train, n_test):
 4    rng=np.random.default_rng(seed); p=3
 5    w=np.array([[.8,-.4,.3],[-.3,.7,.5]],np.float32)
 6    def make(n):
 7        s=rng.normal(size=(n,p)).astype(np.float32)
 8        target=np.tanh(s@w.T+.12*rng.normal(size=(n,2))).astype(np.float32)
 9        a=rng.uniform(-1,1,size=(n,2)).astype(np.float32); b,c=.75,-.25
10        q1=-(a[:,0]-target[:,0])**2+b*a[:,0]*a[:,1]
11        q2=-(a[:,1]-target[:,1])**2+c*a[:,0]*a[:,1]
12        return np.concatenate([s,a],1),np.stack([q1,q2],1).astype(np.float32)
13    xtr,ytr=make(n_train); xte,yte=make(n_test)
14    return {'xtr':xtr,'ytr':ytr,'xte':xte,'yte':yte,'task':'regression','metric':'mse','out_dim':2}