agent.py
def update(self, replay_buffer: rb.ReplayBuffer, batch_size: int, rng: jnp.ndarray) -> None:
        """
            Sample batch of transitions and update both the policy and critic networks.
            As this function contains a conditional function, periodically updating the actor, we do not jit compile it.
        """
        self.updates += 1

        # Provide each element an independent rng sample.
        replay_rand, critic_rand = jax.random.split(rng)

        state, action, next_state, reward, not_done = replay_buffer.sample(batch_size, replay_rand)

        self.critic_params, self.critic_opt_state = self.update_critic(self.critic_params, self.target_critic_params,
                                                                       self.target_actor_params, self.critic_opt_state,
                                                                       state, action, next_state, reward, not_done,
                                                                       critic_rand)

        if self.updates % self.policy_freq == 0:
            self.actor_params, self.actor_opt_state = self.update_actor(self.actor_params, self.critic_params,
                                                                        self.actor_opt_state, state)

            self.target_actor_params = soft_update(self.target_actor_params, self.actor_params)
            self.target_critic_params = soft_update(self.target_critic_params, self.critic_params)

            Sample batch of transitions and update both the policy and critic networks.
            As this function contains a conditional function, periodically updating the actor, we do not jit compile it.
        
Sample batch of transitions and update both the policy and critic networks.
As this function contains a conditional function, periodically updating the actor, we do not jit compile it.
Sample batch of transitions and update both the policy and critic networks . As this function contains a conditional function periodically updating the actor we do not jit compile it .
def update(self, replay_buffer: rb.ReplayBuffer, batch_size: int, rng: jnp.ndarray) -> None:
        self.updates += 1
        replay_rand, critic_rand = jax.random.split(rng)
        state, action, next_state, reward, not_done = replay_buffer.sample(batch_size, replay_rand)
        self.critic_params, self.critic_opt_state = self.update_critic(self.critic_params, self.target_critic_params,
                                                                       self.target_actor_params, self.critic_opt_state,
                                                                       state, action, next_state, reward, not_done,
                                                                       critic_rand)
        if self.updates % self.policy_freq == 0:
            self.actor_params, self.actor_opt_state = self.update_actor(self.actor_params, self.critic_params,
                                                                        self.actor_opt_state, state)
            self.target_actor_params = soft_update(self.target_actor_params, self.actor_params)
            self.target_critic_params = soft_update(self.target_critic_params, self.critic_params)
def update ( self , replay_buffer : rb . ReplayBuffer , batch_size : int , rng : jnp . ndarray ) -> None : self . updates += 1 replay_rand , critic_rand = jax . random . split ( rng ) state , action , next_state , reward , not_done = replay_buffer . sample ( batch_size , replay_rand ) self . critic_params , self . critic_opt_state = self . update_critic ( self . critic_params , self . target_critic_params , self . target_actor_params , self . critic_opt_state , state , action , next_state , reward , not_done , critic_rand ) if self . updates % self . policy_freq == 0 : self . actor_params , self . actor_opt_state = self . update_actor ( self . actor_params , self . critic_params , self . actor_opt_state , state ) self . target_actor_params = soft_update ( self . target_actor_params , self . actor_params ) self . target_critic_params = soft_update ( self . target_critic_params , self . critic_params )
Sample batch of transitions and update both the policy and critic networks.
Sample batch of transitions and update both the policy and critic networks .
"""
            Sample batch of transitions and update both the policy and critic networks.
            As this function contains a conditional function, periodically updating the actor, we do not jit compile it.
        """ # Provide each element an independent rng sample.
self
replay_buffer
rb.ReplayBuffer
batch_size
int
rng
jnp.ndarray