1import numpy as np
 2
 3META = {
 4    'name': 'poisson_action_value',
 5    'domain': 'dynamics/control',
 6    'description': 'Continuous two-dimensional action-value regression with smooth isolated quadratic optima.'
 7}
 8
 9def get_dataset(seed, n_train, n_test):
10    def sample(n, rs):
11        state = rs.uniform(-1.0, 1.0, size=(int(n), 2)).astype(np.float32)
12        optimum = np.stack([
13            0.65 * state[:, 0] + 0.15 * np.sin(2.0 * state[:, 1]),
14            -0.55 * state[:, 1] + 0.10 * np.cos(2.0 * state[:, 0]),
15        ], axis=1).astype(np.float32)
16        action = rs.uniform(-1.0, 1.0, size=(int(n), 2)).astype(np.float32)
17        q = -np.sum((action - optimum) ** 2, axis=1, keepdims=True)
18        return np.concatenate([state, action], axis=1), q
19    xtr, ytr = sample(n_train, np.random.RandomState(int(seed)))
20    xte, yte = sample(n_test, np.random.RandomState(int(seed) + 100003))
21    return {'xtr': xtr, 'ytr': ytr, 'xte': xte, 'yte': yte,
22            'task': 'regression', 'metric': 'mse', 'out_dim': 1}