File size: 1,819 Bytes
7985e7e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | import numpy as np
from .d4rl import load_environment
def compose(*fns):
def _fn(x):
for fn in fns:
x = fn(x)
return x
return _fn
def get_preprocess_fn(fn_names, env):
fns = [eval(name)(env) for name in fn_names]
return compose(*fns)
def get_policy_preprocess_fn(fn_names):
fns = [eval(name) for name in fn_names]
return compose(*fns)
def arctanh_actions(*args, **kwargs):
epsilon = 1e-4
def _fn(dataset):
actions = dataset["actions"]
assert (
actions.min() >= -1 and actions.max() <= 1
), f"applying arctanh to actions in range [{actions.min()}, {actions.max()}]"
actions = np.clip(actions, -1 + epsilon, 1 - epsilon)
dataset["actions"] = np.arctanh(actions)
return dataset
return _fn
def add_deltas(env):
def _fn(dataset):
deltas = dataset["next_observations"] - dataset["observations"]
dataset["deltas"] = deltas
return dataset
return _fn
def maze2d_set_terminals(env):
env = load_environment(env) if type(env) == str else env
goal = np.array(env._target)
threshold = 0.5
def _fn(dataset):
xy = dataset["observations"][:, :2]
distances = np.linalg.norm(xy - goal, axis=-1)
at_goal = distances < threshold
timeouts = np.zeros_like(dataset["timeouts"])
timeouts[:-1] = at_goal[:-1] * ~at_goal[1:]
timeout_steps = np.where(timeouts)[0]
path_lengths = timeout_steps[1:] - timeout_steps[:-1]
print(
f"[ utils/preprocessing ] Segmented {env.name} | {len(path_lengths)} paths | "
f"min length: {path_lengths.min()} | max length: {path_lengths.max()}"
)
dataset["timeouts"] = timeouts
return dataset
return _fn
|