def ClipFraction(dist_inputs, actions, old_log_probs):
"""Probability Ratio Mean from the PPO algorithm."""
probs_ratio = rl_layers.ProbsRatio(
dist_inputs,
actions,
old_log_probs,
log_prob_fun=self._policy_dist.log_prob)
return jnp.mean(jnp.abs(probs_ratio - 1) > self._epsilon)
def forward(self, inputs, weights):
gamma, beta, epsilon_l = weights
epsilon = self._init_epsilon
if epsilon_l is not base.EMPTY_WEIGHTS:
epsilon += np.abs(epsilon_l[0])
# Omit B and C
axis = tuple(range(1, len(np.shape(inputs)) - 1))
# (B, 1, 1, C)
nu2 = np.mean(inputs**2, axis=axis, keepdims=True)
# (B, W, H, C)
xhat = inputs / np.sqrt(nu2 + epsilon)
return gamma * xhat + beta
def SaturationCost(x, limit=0.9): return np.minimum(0, np.abs(x) - limit)