agent.py
def update_critic(
            self,
            critic_params: hk.Params,
            target_critic_params: hk.Params,
            target_actor_params: hk.Params,
            critic_opt_state: OptState,
            state: np.ndarray,
            action: np.ndarray,
            next_state: np.ndarray,
            reward: np.ndarray,
            not_done: np.ndarray,
            rng: jnp.ndarray
    ) -> Tuple[hk.Params, OptState]:
        """Learning rule (stochastic gradient descent)."""
        _, gradient = jax.value_and_grad(self.critic_loss)(critic_params, target_critic_params, target_actor_params,
                                                           state, action, next_state, reward, not_done, rng)
        updates, opt_state = self.critic_opt_update(gradient, critic_opt_state)
        new_params = optix.apply_updates(critic_params, updates)
        return new_params, opt_state
Learning rule (stochastic gradient descent).
Learning rule (stochastic gradient descent).
Learning rule ( stochastic gradient descent ) .
def update_critic(
            self,
            critic_params: hk.Params,
            target_critic_params: hk.Params,
            target_actor_params: hk.Params,
            critic_opt_state: OptState,
            state: np.ndarray,
            action: np.ndarray,
            next_state: np.ndarray,
            reward: np.ndarray,
            not_done: np.ndarray,
            rng: jnp.ndarray
    ) -> Tuple[hk.Params, OptState]:
        _, gradient = jax.value_and_grad(self.critic_loss)(critic_params, target_critic_params, target_actor_params,
                                                           state, action, next_state, reward, not_done, rng)
        updates, opt_state = self.critic_opt_update(gradient, critic_opt_state)
        new_params = optix.apply_updates(critic_params, updates)
        return new_params, opt_state
def update_critic ( self , critic_params : hk . Params , target_critic_params : hk . Params , target_actor_params : hk . Params , critic_opt_state : OptState , state : np . ndarray , action : np . ndarray , next_state : np . ndarray , reward : np . ndarray , not_done : np . ndarray , rng : jnp . ndarray ) -> Tuple [ hk . Params , OptState ] : _ , gradient = jax . value_and_grad ( self . critic_loss ) ( critic_params , target_critic_params , target_actor_params , state , action , next_state , reward , not_done , rng ) updates , opt_state = self . critic_opt_update ( gradient , critic_opt_state ) new_params = optix . apply_updates ( critic_params , updates ) return new_params , opt_state
Learning rule (stochastic gradient descent).
Learning rule ( stochastic gradient descent ) .
"""Learning rule (stochastic gradient descent)."""
self
critic_params
hk.Params
target_critic_params
hk.Params
target_actor_params
hk.Params
critic_opt_state
OptState
state
np.ndarray
action
np.ndarray
next_state
np.ndarray
reward
np.ndarray
not_done
np.ndarray
rng
jnp.ndarray