agent.py
def critic_loss(
            self,
            critic_params: hk.Params,
            target_critic_params: hk.Params,
            target_actor_params: hk.Params,
            state: np.ndarray,
            action: np.ndarray,
            next_state: np.ndarray,
            reward: np.ndarray,
            not_done: np.ndarray,
            rng: jnp.ndarray
    ) -> jnp.DeviceArray:
        """
            TD3 adds truncated Gaussian noise to the policy while training the critic.
            Can be seen as a form of 'Exploration Consciousness' https://arxiv.org/abs/1812.05551 or simply as a
            regularization scheme.
            As this helps stabilize the critic, we also use this for the DDPG update rule.
        """
        noise = (
                jax.random.normal(rng, shape=action.shape) * self.policy_noise
        ).clip(-self.noise_clip, self.noise_clip)

        # Make sure the noisy action is within the valid bounds.
        next_action = (
                self.actor.apply(target_actor_params, next_state) + noise
        ).clip(-self.max_action, self.max_action)

        next_q_1, next_q_2 = self.critic.apply(target_critic_params, jnp.concatenate((next_state, next_action), 1))
        if self.td3_update:
            next_q = jax.lax.min(next_q_1, next_q_2)
        else:
            # Since the actor uses Q_1 for training, setting this as the target for the critic updates is sufficient to
            # obtain an equivalent update.
            next_q = next_q_1
        # Cut the gradient from flowing through the target critic. This is more efficient, computationally.
        target_q = jax.lax.stop_gradient(reward + self.discount * next_q * not_done)
        q_1, q_2 = self.critic.apply(critic_params, jnp.concatenate((state, action), 1))

        return jnp.mean(rlax.l2_loss(q_1, target_q) + rlax.l2_loss(q_2, target_q))

            TD3 adds truncated Gaussian noise to the policy while training the critic.
            Can be seen as a form of 'Exploration Consciousness' https://arxiv.org/abs/1812.05551 or simply as a
            regularization scheme.
            As this helps stabilize the critic, we also use this for the DDPG update rule.
        
TD3 adds truncated Gaussian noise to the policy while training the critic.
TD3 adds truncated Gaussian noise to the policy while training the critic .
def critic_loss(
            self,
            critic_params: hk.Params,
            target_critic_params: hk.Params,
            target_actor_params: hk.Params,
            state: np.ndarray,
            action: np.ndarray,
            next_state: np.ndarray,
            reward: np.ndarray,
            not_done: np.ndarray,
            rng: jnp.ndarray
    ) -> jnp.DeviceArray:
        noise = (
                jax.random.normal(rng, shape=action.shape) * self.policy_noise
        ).clip(-self.noise_clip, self.noise_clip)
        next_action = (
                self.actor.apply(target_actor_params, next_state) + noise
        ).clip(-self.max_action, self.max_action)
        next_q_1, next_q_2 = self.critic.apply(target_critic_params, jnp.concatenate((next_state, next_action), 1))
        if self.td3_update:
            next_q = jax.lax.min(next_q_1, next_q_2)
        else:
            next_q = next_q_1
        target_q = jax.lax.stop_gradient(reward + self.discount * next_q * not_done)
        q_1, q_2 = self.critic.apply(critic_params, jnp.concatenate((state, action), 1))
        return jnp.mean(rlax.l2_loss(q_1, target_q) + rlax.l2_loss(q_2, target_q))
def critic_loss ( self , critic_params : hk . Params , target_critic_params : hk . Params , target_actor_params : hk . Params , state : np . ndarray , action : np . ndarray , next_state : np . ndarray , reward : np . ndarray , not_done : np . ndarray , rng : jnp . ndarray ) -> jnp . DeviceArray : noise = ( jax . random . normal ( rng , shape = action . shape ) * self . policy_noise ) . clip ( - self . noise_clip , self . noise_clip ) next_action = ( self . actor . apply ( target_actor_params , next_state ) + noise ) . clip ( - self . max_action , self . max_action ) next_q_1 , next_q_2 = self . critic . apply ( target_critic_params , jnp . concatenate ( ( next_state , next_action ) , 1 ) ) if self . td3_update : next_q = jax . lax . min ( next_q_1 , next_q_2 ) else : next_q = next_q_1 target_q = jax . lax . stop_gradient ( reward + self . discount * next_q * not_done ) q_1 , q_2 = self . critic . apply ( critic_params , jnp . concatenate ( ( state , action ) , 1 ) ) return jnp . mean ( rlax . l2_loss ( q_1 , target_q ) + rlax . l2_loss ( q_2 , target_q ) )
TD3 adds truncated Gaussian noise to the policy while training the critic.
TD3 adds truncated Gaussian noise to the policy while training the critic .
"""
            TD3 adds truncated Gaussian noise to the policy while training the critic.
            Can be seen as a form of 'Exploration Consciousness' https://arxiv.org/abs/1812.05551 or simply as a
            regularization scheme.
            As this helps stabilize the critic, we also use this for the DDPG update rule.
        """ # Make sure the noisy action is within the valid bounds. # Since the actor uses Q_1 for training, setting this as the target for the critic updates is sufficient to # obtain an equivalent update. # Cut the gradient from flowing through the target critic. This is more efficient, computationally.
self
critic_params
hk.Params
target_critic_params
hk.Params
target_actor_params
hk.Params
state
np.ndarray
action
np.ndarray
next_state
np.ndarray
reward
np.ndarray
not_done
np.ndarray
rng
jnp.ndarray