def dot(self, x, y):
if self.dim == 2:
rxs, thxs = x[:, 0], x[:, 1]
rys, thys = y[:, 0], y[:, 1]
return 4 * np.arctanh(rxs) * np.arctanh(rys) * np.cos(thxs - thys)
else:
pass
def evaluate_gaussian_and_tanh_log_prob(
mean: jnp.ndarray,
log_std: jnp.ndarray,
action: jnp.ndarray,
) -> jnp.ndarray:
"""
Calculate log probabilities of gaussian distributions and tanh transformation given samples.
"""
noise = (jnp.arctanh(action) - mean) / (jnp.exp(log_std) + 1e-8)
return gaussian_and_tanh_log_prob(log_std, noise,
action).sum(axis=1, keepdims=True)
def _inverse(self, y):
# inverse stick-breaking
remainder = 1 - jnp.cumsum(jnp.abs(y[..., :-1]), axis=-1)
pad_width = [(0, 0)] * (y.ndim - 1) + [(1, 0)]
remainder = jnp.pad(remainder, pad_width, mode="constant", constant_values=1.0)
finfo = jnp.finfo(y.dtype)
remainder = jnp.clip(remainder, a_min=finfo.tiny)
t = y / remainder
# inverse of tanh
t = jnp.clip(t, a_min=-1 + finfo.eps, a_max=1 - finfo.eps)
return jnp.arctanh(t)
def _inverse(self, y):
# inverse stick-breaking
z1m_cumprod = 1 - jnp.cumsum(y * y, axis=-1)
pad_width = [(0, 0)] * y.ndim
pad_width[-1] = (1, 0)
z1m_cumprod_shifted = jnp.pad(z1m_cumprod[..., :-1],
pad_width,
mode="constant",
constant_values=1.0)
t = matrix_to_tril_vec(y, diagonal=-1) / jnp.sqrt(
matrix_to_tril_vec(z1m_cumprod_shifted, diagonal=-1))
# inverse of tanh
return jnp.arctanh(t)
def fn(x):
return jnp.arctanh(x) * hk.get_parameter(
'p', [], init=hk.initializers.Constant(0.))
theoretical_variance = 9.64532433084641
theoretical_sigma = np.sqrt(theoretical_variance / n_samples)
bias_dict = dict()
for dt in dts:
solver.dt = dt
solution = solver.solve_many(problem,n_samples,seed=0)
terminal_x = solution['solution_values'][:,-1,0]
expected_x = float(jnp.mean(terminal_x))
std_mean_x = float(jnp.std(terminal_x) / jnp.sqrt(n_samples))
measured_bias = expected_x-theoretical_mean
sigma_bias = std_mean_x
true_x = jnp.tanh( solution['wiener_values'][:,-1,0] + jnp.arctanh(problem.x0))
print(f'{jnp.mean(jnp.abs(terminal_x - true_x))=}')
print(f'{jnp.mean(true_x)-theoretical_mean=}')
print(f'{int(1/dt)=}')
print(f'{measured_bias=}')
print(f'{sigma_bias=}')
print(f'{measured_bias/sigma_bias=}')
bias_dict[int(1/dt)] = (measured_bias,sigma_bias)
for (k,(mu,sigma)) in bias_dict.items():
print('{'+f'{k},{mu},{sigma}'+'},')
def inverse_and_log_det(self, y: Array) -> Tuple[Array, Array]:
"""Computes x = f^{-1}(y) and log|det J(f^{-1})(y)|."""
x = jnp.arctanh(y)
return x, -self.forward_log_det_jacobian(x)
def arctanh(x): if isinstance(x, JaxArray): x = x.value return JaxArray(jnp.arctanh(x))
def arctanh(a: Numeric):
return jnp.arctanh(a)
def inv_transform(a, action_spec):
min_vals, max_vals = minmaxvals(a, action_spec)
scale = (max_vals - min_vals) * 0.5
actions_tanh = (a - min_vals) / scale - 1.
return jnp.arctanh(actions_tanh)
def arctanh(self: TensorType) -> TensorType:
return type(self)(np.arctanh(self.raw))
def inverse(self, value):
return np.arctanh(value)
def _inverse(self, y):
# We perform clipping in the _inverse function, as is done in TF-Agents.
y = jnp.where(jnp.less_equal(jnp.abs(y), 1.),
tf.clip(y, -0.99999997, 0.99999997), y)
return jnp.arctanh(y)
def deconstrain(v, a, b):
return jnp.arctanh(jnp.clip((v - a) * 2. / (b - a) - 1., -0.999, 0.999))
def norm(self, x):
if self.dim == 2:
return 2 * np.arctanh(x[:, 0])
else:
euc_norms = np.sqrt(np.sum(x ** 2, axis=-1))
return 2 * np.arctanh(euc_norms)
argmin = utils.copy_docstring(
tf.math.argmin,
lambda input, axis=None, output_type=tf.int64, name=None: ( # pylint: disable=g-long-lambda
np.argmin(input, axis=0 if axis is None else _astuple(axis)).astype(
utils.numpy_dtype(output_type))))
asin = utils.copy_docstring(tf.math.asin, lambda x, name=None: np.arcsin(x))
asinh = utils.copy_docstring(tf.math.asinh, lambda x, name=None: np.arcsinh(x))
atan = utils.copy_docstring(tf.math.atan, lambda x, name=None: np.arctan(x))
atan2 = utils.copy_docstring(tf.math.atan2,
lambda y, x, name=None: np.arctan2(y, x))
atanh = utils.copy_docstring(tf.math.atanh, lambda x, name=None: np.arctanh(x))
bessel_i0 = utils.copy_docstring(tf.math.bessel_i0,
lambda x, name=None: scipy_special.i0(x))
bessel_i0e = utils.copy_docstring(tf.math.bessel_i0e,
lambda x, name=None: scipy_special.i0e(x))
bessel_i1 = utils.copy_docstring(tf.math.bessel_i1,
lambda x, name=None: scipy_special.i1(x))
bessel_i1e = utils.copy_docstring(tf.math.bessel_i1e,
lambda x, name=None: scipy_special.i1e(x))
betainc = utils.copy_docstring(
tf.math.betainc, lambda a, b, x, name=None: scipy_special.betainc(a, b, x))
