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