1import numpy as np
 2import torch
 3
 4META = {
 5    "name": "incidence_actions",
 6    "domain": "structured_action_classification",
 7    "description": "Variable-size atomic entities and legal composite actions represented by an instance-specific binary incidence matrix; labels are noisy utility-maximizing legal actions."
 8}
 9D = 4
10MAX_ATOMS = 24
11MAX_ACTIONS = 48
12
13def get_dataset(seed, n_train, n_test):
14    rng = np.random.default_rng(int(seed))
15    true_w = np.array([1.2, -0.9, 0.65, 0.35], dtype=np.float32)
16    def make(n):
17        X = np.zeros((n, MAX_ATOMS * D + MAX_ACTIONS * MAX_ATOMS), dtype=np.float32)
18        y = np.zeros(n, dtype=np.int64)
19        for b in range(n):
20            na = int(rng.integers(5, MAX_ATOMS + 1))
21            pa = int(rng.integers(5, MAX_ACTIONS + 1))
22            atoms = rng.normal(size=(na, D)).astype(np.float32)
23            A = np.zeros((pa, na), dtype=np.float32)
24            for p in range(pa):
25                k = int(rng.integers(1, min(6, na) + 1))
26                A[p, rng.choice(na, size=k, replace=False)] = 1.0
27            utilities = A @ (atoms @ true_w) + 0.08 * rng.normal(size=pa)
28            y[b] = int(np.argmax(utilities))
29            X[b, :na * D] = atoms.reshape(-1)
30            packed = np.pad(A, ((0, 0), (0, MAX_ATOMS-na))).reshape(-1)
31            X[b, MAX_ATOMS * D:MAX_ATOMS * D + pa * MAX_ATOMS] = packed
32        return X, y
33    xtr, ytr = make(n_train)
34    xte, yte = make(n_test)
35    return {"xtr": torch.from_numpy(xtr), "ytr": torch.from_numpy(ytr),
36            "xte": torch.from_numpy(xte), "yte": torch.from_numpy(yte),
37            "input_shape": (xtr.shape[1],), "out_dim": MAX_ACTIONS,
38            "task": "classification", "metric": "cross_entropy"}