1import numpy as np
 2
 3META = {
 4    "name": "obb_clearance_trajectory",
 5    "domain": "geometric_collision_avoidance",
 6    "description": "Eight-step ego trajectory regression with an observed oriented obstacle box."
 7}
 8
 9def get_dataset(seed, n_train, n_test):
10    def make(n, s):
11        r = np.random.RandomState(s)
12        start_y = r.uniform(-.12, .12, n)
13        obs_x = r.uniform(1.8, 2.2, n)
14        obs_y = r.uniform(-.18, .18, n)
15        obs_yaw = r.uniform(-.25, .25, n)
16        x = np.stack([start_y, obs_x, obs_y, obs_yaw], 1).astype(np.float32)
17        t = np.linspace(.5, 4., 8, dtype=np.float32)[None, :]
18        yy = start_y[:, None] + (obs_y-start_y)[:, None] * t/4
19        y = np.stack([np.broadcast_to(t, (n,8)), yy], -1)
20        y += r.normal(0, .015, y.shape).astype(np.float32)
21        return x, y.reshape(n,16).astype(np.float32)
22    xtr,ytr = make(n_train, seed)
23    xte,yte = make(n_test, seed+5000)
24    return {"xtr":xtr,"ytr":ytr,"xte":xte,"yte":yte,
25            "task":"regression","metric":"mse","out_dim":16}