def toy_feature_vec_func(s):
"""
:param s: self, (T, 2)
:param o: other, (T, 2)
:return: features, (T, Df, 2)
"""
corners = torch.tensor([[0, 0], [0, 8], [10, 0], [10, 8]], dtype=torch.float64)
return feature_direction_vec(s, corners)
x_grid_gap = (arena_xmax - arena_xmin) / 3
y_grid_gap = (arena_ymax - arena_ymin) / 3
x_grids = [arena_xmin + i * x_grid_gap for i in range(4)]
y_grids = [arena_ymin + i * y_grid_gap for i in range(4)]
torch.manual_seed(0)
np.random.seed(0)
D = 4
K = 2
M = 0
Df = 4
f_corner_vec_func = lambda s: feature_direction_vec(s, CORNERS)
bounds = np.array([[arena_xmin, arena_xmax], [arena_ymin, arena_ymax],
[arena_xmin, arena_xmax], [arena_ymin, arena_ymax]])
tran = GridTransformation(K=K,
D=D,
x_grids=x_grids,
y_grids=y_grids,
unit_transformation="direction",
Df=Df,
feature_vec_func=f_corner_vec_func,
acc_factor=10)
obs = ARTruncatedNormalObservation(K=K,
D=D,
