def __call__( self, network: networks_lib.FeedForwardNetwork, params: networks_lib.Params, target_params: networks_lib.Params, batch: reverb.ReplaySample, key: networks_lib.PRNGKey, ) -> Tuple[jnp.DeviceArray, learning_lib.LossExtra]: """Calculate a loss on a single batch of data.""" del key transitions: types.Transition = batch.data # Forward pass. q_tm1 = network.apply(params, transitions.observation) q_t = network.apply(target_params, transitions.next_observation) d_t = (transitions.discount * self.discount).astype(jnp.float32) # Compute Q-learning TD-error. batch_error = jax.vmap(rlax.q_learning) td_error = batch_error(q_tm1, transitions.action, transitions.reward, d_t, q_t) td_error = 0.5 * jnp.square(td_error) def select(qtm1, action): return qtm1[action] q_regularizer = jax.vmap(select)(q_tm1, transitions.action) loss = self.regularizer_coeff * jnp.mean(q_regularizer) + jnp.mean( td_error) extra = learning_lib.LossExtra(metrics={}) return loss, extra
def __call__( self, network: networks_lib.FeedForwardNetwork, params: networks_lib.Params, target_params: networks_lib.Params, batch: reverb.ReplaySample, key: networks_lib.PRNGKey, ) -> Tuple[jnp.DeviceArray, learning_lib.LossExtra]: """Calculate a loss on a single batch of data.""" del key transitions: types.Transition = batch.data # Forward pass. q_tm1 = network.apply(params, transitions.observation) q_t = network.apply(target_params, transitions.next_observation) # Cast and clip rewards. d_t = (transitions.discount * self.discount).astype(jnp.float32) r_t = jnp.clip(transitions.reward, -self.max_abs_reward, self.max_abs_reward).astype(jnp.float32) # Compute Q-learning TD-error. batch_error = jax.vmap(rlax.q_learning) td_error = batch_error(q_tm1, transitions.action, r_t, d_t, q_t) batch_loss = jnp.square(td_error) loss = jnp.mean(batch_loss) extra = learning_lib.LossExtra(metrics={}) return loss, extra
def default_behavior_policy(network: networks_lib.FeedForwardNetwork, epsilon: float, params: networks_lib.Params, key: networks_lib.PRNGKey, observation: networks_lib.Observation): """Returns an action for the given observation.""" action_values = network.apply(params, observation) actions = rlax.epsilon_greedy(epsilon).sample(key, action_values) return actions.astype(jnp.int32)
def __call__( self, network: networks_lib.FeedForwardNetwork, params: networks_lib.Params, target_params: networks_lib.Params, batch: reverb.ReplaySample, key: networks_lib.PRNGKey, ) -> Tuple[jnp.DeviceArray, learning_lib.LossExtra]: """Calculate a loss on a single batch of data.""" del key transitions: types.Transition = batch.data # Forward pass. q_online_s = network.apply(params, transitions.observation) action_one_hot = jax.nn.one_hot(transitions.action, q_online_s.shape[-1]) q_online_sa = jnp.sum(action_one_hot * q_online_s, axis=-1) q_target_s = network.apply(target_params, transitions.observation) q_target_next = network.apply(target_params, transitions.next_observation) # Cast and clip rewards. d_t = (transitions.discount * self.discount).astype(jnp.float32) r_t = jnp.clip(transitions.reward, -self.max_abs_reward, self.max_abs_reward).astype(jnp.float32) # Munchausen term : tau * log_pi(a|s) munchausen_term = self.entropy_temperature * jax.nn.log_softmax( q_target_s / self.entropy_temperature, axis=-1) munchausen_term_a = jnp.sum(action_one_hot * munchausen_term, axis=-1) munchausen_term_a = jnp.clip(munchausen_term_a, a_min=self.clip_value_min, a_max=0.) # Soft Bellman operator applied to q next_v = self.entropy_temperature * jax.nn.logsumexp( q_target_next / self.entropy_temperature, axis=-1) target_q = jax.lax.stop_gradient(r_t + self.munchausen_coefficient * munchausen_term_a + d_t * next_v) batch_loss = rlax.huber_loss(target_q - q_online_sa, self.huber_loss_parameter) loss = jnp.mean(batch_loss) extra = learning_lib.LossExtra(metrics={}) return loss, extra
def __call__( self, network: networks_lib.FeedForwardNetwork, params: networks_lib.Params, target_params: networks_lib.Params, batch: reverb.ReplaySample, key: networks_lib.PRNGKey, ) -> Tuple[jnp.DeviceArray, learning_lib.LossExtra]: """Calculate a loss on a single batch of data.""" del key transitions: types.Transition = batch.data keys, probs, *_ = batch.info # Forward pass. _, logits_tm1, atoms_tm1 = network.apply(params, transitions.observation) _, logits_t, atoms_t = network.apply(target_params, transitions.next_observation) q_t_selector, _, _ = network.apply(params, transitions.next_observation) # Cast and clip rewards. d_t = (transitions.discount * self.discount).astype(jnp.float32) r_t = jnp.clip(transitions.reward, -self.max_abs_reward, self.max_abs_reward).astype(jnp.float32) # Compute categorical double Q-learning loss. batch_loss_fn = jax.vmap(rlax.categorical_double_q_learning, in_axes=(None, 0, 0, 0, 0, None, 0, 0)) batch_loss = batch_loss_fn(atoms_tm1, logits_tm1, transitions.action, r_t, d_t, atoms_t, logits_t, q_t_selector) # Importance weighting. importance_weights = (1. / probs).astype(jnp.float32) importance_weights **= self.importance_sampling_exponent importance_weights /= jnp.max(importance_weights) # Reweight. loss = jnp.mean(importance_weights * batch_loss) # [] reverb_update = learning_lib.ReverbUpdate( keys=keys, priorities=jnp.abs(batch_loss).astype(jnp.float64)) extra = learning_lib.LossExtra(metrics={}, reverb_update=reverb_update) return loss, extra
def __call__( self, network: networks_lib.FeedForwardNetwork, params: networks_lib.Params, target_params: networks_lib.Params, batch: reverb.ReplaySample, key: jnp.DeviceArray, ) -> Tuple[jnp.DeviceArray, learning_lib.LossExtra]: """Calculate a loss on a single batch of data.""" del key transitions: types.Transition = batch.data keys, probs, *_ = batch.info # Forward pass. q_tm1 = network.apply(params, transitions.observation) q_t_value = network.apply(target_params, transitions.next_observation) q_t_selector = network.apply(params, transitions.next_observation) # Cast and clip rewards. d_t = (transitions.discount * self.discount).astype(jnp.float32) r_t = jnp.clip(transitions.reward, -self.max_abs_reward, self.max_abs_reward).astype(jnp.float32) # Compute double Q-learning n-step TD-error. batch_error = jax.vmap(rlax.double_q_learning) td_error = batch_error(q_tm1, transitions.action, r_t, d_t, q_t_value, q_t_selector) batch_loss = rlax.huber_loss(td_error, self.huber_loss_parameter) # Importance weighting. importance_weights = (1. / probs).astype(jnp.float32) importance_weights **= self.importance_sampling_exponent importance_weights /= jnp.max(importance_weights) # Reweight. loss = jnp.mean(importance_weights * batch_loss) # [] reverb_update = learning_lib.ReverbUpdate( keys=keys, priorities=jnp.abs(td_error).astype(jnp.float64)) extra = learning_lib.LossExtra(metrics={}, reverb_update=reverb_update) return loss, extra
def __call__( self, network: networks_lib.FeedForwardNetwork, params: networks_lib.Params, target_params: networks_lib.Params, batch: reverb.ReplaySample, key: networks_lib.PRNGKey, ) -> Tuple[jnp.DeviceArray, learning_lib.LossExtra]: """Calculate a loss on a single batch of data.""" del key transitions: types.Transition = batch.data dist_q_tm1 = network.apply(params, transitions.observation)['q_dist'] dist_q_target_t = network.apply(target_params, transitions.next_observation)['q_dist'] # Swap distribution and action dimension, since # rlax.quantile_q_learning expects it that way. dist_q_tm1 = jnp.swapaxes(dist_q_tm1, 1, 2) dist_q_target_t = jnp.swapaxes(dist_q_target_t, 1, 2) quantiles = ((jnp.arange(self.num_atoms, dtype=jnp.float32) + 0.5) / self.num_atoms) batch_quantile_q_learning = jax.vmap(rlax.quantile_q_learning, in_axes=(0, None, 0, 0, 0, 0, 0, None)) losses = batch_quantile_q_learning( dist_q_tm1, quantiles, transitions.action, transitions.reward, transitions.discount, dist_q_target_t, # No double Q-learning here. dist_q_target_t, self.huber_param, ) loss = jnp.mean(losses) extra = learning_lib.LossExtra(metrics={'mean_loss': loss}) return loss, extra
def __init__(self, spec: specs.EnvironmentSpec, networks: networks_lib.FeedForwardNetwork, rng: networks_lib.PRNGKey, config: ars_config.ARSConfig, iterator: Iterator[reverb.ReplaySample], counter: Optional[counting.Counter] = None, logger: Optional[loggers.Logger] = None): self._config = config self._lock = threading.Lock() # General learner book-keeping and loggers. self._counter = counter or counting.Counter() self._logger = logger or loggers.make_default_logger( 'learner', asynchronous=True, serialize_fn=utils.fetch_devicearray, steps_key=self._counter.get_steps_key()) # Iterator on demonstration transitions. self._iterator = iterator if self._config.normalize_observations: normalizer_params = running_statistics.init_state( spec.observations) self._normalizer_update_fn = running_statistics.update else: normalizer_params = () self._normalizer_update_fn = lambda a, b: a rng1, rng2, tmp = jax.random.split(rng, 3) # Create initial state. self._training_state = TrainingState( key=rng1, policy_params=networks.init(tmp), normalizer_params=normalizer_params, training_iteration=0) self._evaluation_state = EvaluationState( key=rng2, evaluation_queue=collections.deque(), received_results={}, noises=[]) # Do not record timestamps until after the first learning step is done. # This is to avoid including the time it takes for actors to come online and # fill the replay buffer. self._timestamp = None
def __init__(self, network: networks_lib.FeedForwardNetwork, obs_spec: specs.Array, optimizer: optax.GradientTransformation, random_key: networks_lib.PRNGKey, dataset: tf.data.Dataset, loss_fn: LossFn = _sparse_categorical_cross_entropy, counter: counting.Counter = None, logger: loggers.Logger = None): """Initializes the learner.""" def loss(params: networks_lib.Params, sample: reverb.ReplaySample) -> jnp.ndarray: # Pull out the data needed for updates. o_tm1, a_tm1, r_t, d_t, o_t = sample.data del r_t, d_t, o_t logits = network.apply(params, o_tm1) return jnp.mean(loss_fn(a_tm1, logits)) def sgd_step( state: TrainingState, sample: reverb.ReplaySample ) -> Tuple[TrainingState, Dict[str, jnp.DeviceArray]]: """Do a step of SGD.""" grad_fn = jax.value_and_grad(loss) loss_value, gradients = grad_fn(state.params, sample) updates, new_opt_state = optimizer.update(gradients, state.opt_state) new_params = optax.apply_updates(state.params, updates) steps = state.steps + 1 new_state = TrainingState(params=new_params, opt_state=new_opt_state, steps=steps) # Compute the global norm of the gradients for logging. global_gradient_norm = optax.global_norm(gradients) fetches = { 'loss': loss_value, 'gradient_norm': global_gradient_norm } return new_state, fetches self._counter = counter or counting.Counter() self._logger = logger or loggers.TerminalLogger('learner', time_delta=1.) # Get an iterator over the dataset. self._iterator = iter(dataset) # pytype: disable=wrong-arg-types # TODO(b/155086959): Fix type stubs and remove. # Initialise parameters and optimiser state. initial_params = network.init( random_key, utils.add_batch_dim(utils.zeros_like(obs_spec))) initial_opt_state = optimizer.init(initial_params) self._state = TrainingState(params=initial_params, opt_state=initial_opt_state, steps=0) self._sgd_step = jax.jit(sgd_step)
def __init__(self, unroll: networks_lib.FeedForwardNetwork, initial_state: networks_lib.FeedForwardNetwork, batch_size: int, random_key: networks_lib.PRNGKey, burn_in_length: int, discount: float, importance_sampling_exponent: float, max_priority_weight: float, target_update_period: int, iterator: Iterator[reverb.ReplaySample], optimizer: optax.GradientTransformation, bootstrap_n: int = 5, tx_pair: rlax.TxPair = rlax.SIGNED_HYPERBOLIC_PAIR, clip_rewards: bool = False, max_abs_reward: float = 1., use_core_state: bool = True, prefetch_size: int = 2, replay_client: Optional[reverb.Client] = None, counter: Optional[counting.Counter] = None, logger: Optional[loggers.Logger] = None): """Initializes the learner.""" random_key, key_initial_1, key_initial_2 = jax.random.split( random_key, 3) initial_state_params = initial_state.init(key_initial_1, batch_size) initial_state = initial_state.apply(initial_state_params, key_initial_2, batch_size) def loss( params: networks_lib.Params, target_params: networks_lib.Params, key_grad: networks_lib.PRNGKey, sample: reverb.ReplaySample ) -> Tuple[jnp.ndarray, Tuple[jnp.ndarray, jnp.ndarray]]: """Computes mean transformed N-step loss for a batch of sequences.""" # Convert sample data to sequence-major format [T, B, ...]. data = utils.batch_to_sequence(sample.data) # Get core state & warm it up on observations for a burn-in period. if use_core_state: # Replay core state. online_state = jax.tree_map(lambda x: x[0], data.extras['core_state']) else: online_state = initial_state target_state = online_state # Maybe burn the core state in. if burn_in_length: burn_obs = jax.tree_map(lambda x: x[:burn_in_length], data.observation) key_grad, key1, key2 = jax.random.split(key_grad, 3) _, online_state = unroll.apply(params, key1, burn_obs, online_state) _, target_state = unroll.apply(target_params, key2, burn_obs, target_state) # Only get data to learn on from after the end of the burn in period. data = jax.tree_map(lambda seq: seq[burn_in_length:], data) # Unroll on sequences to get online and target Q-Values. key1, key2 = jax.random.split(key_grad) online_q, _ = unroll.apply(params, key1, data.observation, online_state) target_q, _ = unroll.apply(target_params, key2, data.observation, target_state) # Get value-selector actions from online Q-values for double Q-learning. selector_actions = jnp.argmax(online_q, axis=-1) # Preprocess discounts & rewards. discounts = (data.discount * discount).astype(online_q.dtype) rewards = data.reward if clip_rewards: rewards = jnp.clip(rewards, -max_abs_reward, max_abs_reward) rewards = rewards.astype(online_q.dtype) # Get N-step transformed TD error and loss. batch_td_error_fn = jax.vmap(functools.partial( rlax.transformed_n_step_q_learning, n=bootstrap_n, tx_pair=tx_pair), in_axes=1, out_axes=1) # TODO(b/183945808): when this bug is fixed, truncations of actions, # rewards, and discounts will no longer be necessary. batch_td_error = batch_td_error_fn(online_q[:-1], data.action[:-1], target_q[1:], selector_actions[1:], rewards[:-1], discounts[:-1]) batch_loss = 0.5 * jnp.square(batch_td_error).sum(axis=0) # Importance weighting. probs = sample.info.probability importance_weights = (1. / (probs + 1e-6)).astype(online_q.dtype) importance_weights **= importance_sampling_exponent importance_weights /= jnp.max(importance_weights) mean_loss = jnp.mean(importance_weights * batch_loss) # Calculate priorities as a mixture of max and mean sequence errors. abs_td_error = jnp.abs(batch_td_error).astype(online_q.dtype) max_priority = max_priority_weight * jnp.max(abs_td_error, axis=0) mean_priority = (1 - max_priority_weight) * jnp.mean(abs_td_error, axis=0) priorities = (max_priority + mean_priority) return mean_loss, priorities def sgd_step( state: TrainingState, samples: reverb.ReplaySample ) -> Tuple[TrainingState, jnp.ndarray, Dict[str, jnp.ndarray]]: """Performs an update step, averaging over pmap replicas.""" # Compute loss and gradients. grad_fn = jax.value_and_grad(loss, has_aux=True) key, key_grad = jax.random.split(state.random_key) (loss_value, priorities), gradients = grad_fn(state.params, state.target_params, key_grad, samples) # Average gradients over pmap replicas before optimizer update. gradients = jax.lax.pmean(gradients, _PMAP_AXIS_NAME) # Apply optimizer updates. updates, new_opt_state = optimizer.update(gradients, state.opt_state) new_params = optax.apply_updates(state.params, updates) # Periodically update target networks. steps = state.steps + 1 target_params = rlax.periodic_update(new_params, state.target_params, steps, self._target_update_period) new_state = TrainingState(params=new_params, target_params=target_params, opt_state=new_opt_state, steps=steps, random_key=key) return new_state, priorities, {'loss': loss_value} def update_priorities(keys_and_priorities: Tuple[jnp.ndarray, jnp.ndarray]): keys, priorities = keys_and_priorities keys, priorities = tree.map_structure( # Fetch array and combine device and batch dimensions. lambda x: utils.fetch_devicearray(x).reshape( (-1, ) + x.shape[2:]), (keys, priorities)) replay_client.mutate_priorities( # pytype: disable=attribute-error table=adders.DEFAULT_PRIORITY_TABLE, updates=dict(zip(keys, priorities))) # Internalise components, hyperparameters, logger, counter, and methods. self._replay_client = replay_client self._target_update_period = target_update_period self._counter = counter or counting.Counter() self._logger = logger or loggers.make_default_logger( 'learner', asynchronous=True, serialize_fn=utils.fetch_devicearray, time_delta=1.) self._sgd_step = jax.pmap(sgd_step, axis_name=_PMAP_AXIS_NAME) self._async_priority_updater = async_utils.AsyncExecutor( update_priorities) # Initialise and internalise training state (parameters/optimiser state). random_key, key_init = jax.random.split(random_key) initial_params = unroll.init(key_init, initial_state) opt_state = optimizer.init(initial_params) state = TrainingState(params=initial_params, target_params=initial_params, opt_state=opt_state, steps=jnp.array(0), random_key=random_key) # Replicate parameters. self._state = utils.replicate_in_all_devices(state) # Shard multiple inputs with on-device prefetching. # We split samples in two outputs, the keys which need to be kept on-host # since int64 arrays are not supported in TPUs, and the entire sample # separately so it can be sent to the sgd_step method. def split_sample( sample: reverb.ReplaySample) -> utils.PrefetchingSplit: return utils.PrefetchingSplit(host=sample.info.key, device=sample) self._prefetched_iterator = utils.sharded_prefetch( iterator, buffer_size=prefetch_size, num_threads=jax.local_device_count(), split_fn=split_sample)
def __init__(self, policy_network: networks_lib.FeedForwardNetwork, critic_network: networks_lib.FeedForwardNetwork, random_key: networks_lib.PRNGKey, discount: float, target_update_period: int, iterator: Iterator[reverb.ReplaySample], policy_optimizer: Optional[optax.GradientTransformation] = None, critic_optimizer: Optional[optax.GradientTransformation] = None, clipping: bool = True, counter: Optional[counting.Counter] = None, logger: Optional[loggers.Logger] = None, jit: bool = True, num_sgd_steps_per_step: int = 1): def critic_mean( critic_params: networks_lib.Params, observation: types.NestedArray, action: types.NestedArray, ) -> jnp.ndarray: # We add batch dimension to make sure batch concat in critic_network # works correctly. observation = utils.add_batch_dim(observation) action = utils.add_batch_dim(action) # Computes the mean action-value estimate. logits, atoms = critic_network.apply(critic_params, observation, action) logits = utils.squeeze_batch_dim(logits) probabilities = jax.nn.softmax(logits) return jnp.sum(probabilities * atoms, axis=-1) def policy_loss( policy_params: networks_lib.Params, critic_params: networks_lib.Params, o_t: types.NestedArray, ) -> jnp.ndarray: # Computes the discrete policy gradient loss. dpg_a_t = policy_network.apply(policy_params, o_t) grad_critic = jax.vmap( jax.grad(critic_mean, argnums=2), in_axes=(None, 0, 0)) dq_da = grad_critic(critic_params, o_t, dpg_a_t) dqda_clipping = 1. if clipping else None batch_dpg_learning = jax.vmap(rlax.dpg_loss, in_axes=(0, 0, None)) loss = batch_dpg_learning(dpg_a_t, dq_da, dqda_clipping) return jnp.mean(loss) def critic_loss( critic_params: networks_lib.Params, state: TrainingState, transition: types.Transition, ): # Computes the distributional critic loss. q_tm1, atoms_tm1 = critic_network.apply(critic_params, transition.observation, transition.action) a = policy_network.apply(state.target_policy_params, transition.next_observation) q_t, atoms_t = critic_network.apply(state.target_critic_params, transition.next_observation, a) batch_td_learning = jax.vmap( rlax.categorical_td_learning, in_axes=(None, 0, 0, 0, None, 0)) loss = batch_td_learning(atoms_tm1, q_tm1, transition.reward, discount * transition.discount, atoms_t, q_t) return jnp.mean(loss) def sgd_step( state: TrainingState, transitions: types.Transition, ) -> Tuple[TrainingState, Dict[str, jnp.ndarray]]: # TODO(jaslanides): Use a shared forward pass for efficiency. policy_loss_and_grad = jax.value_and_grad(policy_loss) critic_loss_and_grad = jax.value_and_grad(critic_loss) # Compute losses and their gradients. policy_loss_value, policy_gradients = policy_loss_and_grad( state.policy_params, state.critic_params, transitions.next_observation) critic_loss_value, critic_gradients = critic_loss_and_grad( state.critic_params, state, transitions) # Get optimizer updates and state. policy_updates, policy_opt_state = policy_optimizer.update( # pytype: disable=attribute-error policy_gradients, state.policy_opt_state) critic_updates, critic_opt_state = critic_optimizer.update( # pytype: disable=attribute-error critic_gradients, state.critic_opt_state) # Apply optimizer updates to parameters. policy_params = optax.apply_updates(state.policy_params, policy_updates) critic_params = optax.apply_updates(state.critic_params, critic_updates) steps = state.steps + 1 # Periodically update target networks. target_policy_params, target_critic_params = optax.periodic_update( (policy_params, critic_params), (state.target_policy_params, state.target_critic_params), steps, self._target_update_period) new_state = TrainingState( policy_params=policy_params, critic_params=critic_params, target_policy_params=target_policy_params, target_critic_params=target_critic_params, policy_opt_state=policy_opt_state, critic_opt_state=critic_opt_state, steps=steps, ) metrics = { 'policy_loss': policy_loss_value, 'critic_loss': critic_loss_value, } return new_state, metrics # General learner book-keeping and loggers. self._counter = counter or counting.Counter() self._logger = logger or loggers.make_default_logger( 'learner', asynchronous=True, serialize_fn=utils.fetch_devicearray, steps_key=self._counter.get_steps_key()) # Necessary to track when to update target networks. self._target_update_period = target_update_period # Create prefetching dataset iterator. self._iterator = iterator # Maybe use the JIT compiler. sgd_step = utils.process_multiple_batches(sgd_step, num_sgd_steps_per_step) self._sgd_step = jax.jit(sgd_step) if jit else sgd_step # Create the network parameters and copy into the target network parameters. key_policy, key_critic = jax.random.split(random_key) initial_policy_params = policy_network.init(key_policy) initial_critic_params = critic_network.init(key_critic) initial_target_policy_params = initial_policy_params initial_target_critic_params = initial_critic_params # Create optimizers if they aren't given. critic_optimizer = critic_optimizer or optax.adam(1e-4) policy_optimizer = policy_optimizer or optax.adam(1e-4) # Initialize optimizers. initial_policy_opt_state = policy_optimizer.init(initial_policy_params) # pytype: disable=attribute-error initial_critic_opt_state = critic_optimizer.init(initial_critic_params) # pytype: disable=attribute-error # Create initial state. self._state = TrainingState( policy_params=initial_policy_params, target_policy_params=initial_target_policy_params, critic_params=initial_critic_params, target_critic_params=initial_target_critic_params, policy_opt_state=initial_policy_opt_state, critic_opt_state=initial_critic_opt_state, steps=0, ) # Do not record timestamps until after the first learning step is done. # This is to avoid including the time it takes for actors to come online and # fill the replay buffer. self._timestamp = None
def __init__(self, network: networks_lib.FeedForwardNetwork, random_key: networks_lib.PRNGKey, loss_fn: losses.Loss, optimizer: optax.GradientTransformation, prefetching_iterator: Iterator[types.Transition], num_sgd_steps_per_step: int, loss_has_aux: bool = False, logger: Optional[loggers.Logger] = None, counter: Optional[counting.Counter] = None): """Behavior Cloning Learner. Args: network: Networks with signature for apply: (params, obs, is_training, key) -> jnp.ndarray and for init: (rng, is_training) -> params random_key: RNG key. loss_fn: BC loss to use. optimizer: Optax optimizer. prefetching_iterator: A sharded prefetching iterator as outputted from `acme.jax.utils.sharded_prefetch`. Please see the documentation for `sharded_prefetch` for more details. num_sgd_steps_per_step: Number of gradient updates per step. loss_has_aux: Whether the loss function returns auxiliary metrics as a second argument. logger: Logger. counter: Counter. """ def sgd_step( state: TrainingState, transitions: types.Transition, ) -> Tuple[TrainingState, Dict[str, jnp.ndarray]]: loss_and_grad = jax.value_and_grad(loss_fn, argnums=1, has_aux=loss_has_aux) # Compute losses and their gradients. key, key_input = jax.random.split(state.key) loss_result, gradients = loss_and_grad(network.apply, state.policy_params, key_input, transitions) # Combine the gradient across all devices (by taking their mean). gradients = jax.lax.pmean(gradients, axis_name=_PMAP_AXIS_NAME) # Compute and combine metrics across all devices. metrics = _create_loss_metrics(loss_has_aux, loss_result, gradients) metrics = jax.lax.pmean(metrics, axis_name=_PMAP_AXIS_NAME) policy_update, optimizer_state = optimizer.update( gradients, state.optimizer_state, state.policy_params) policy_params = optax.apply_updates(state.policy_params, policy_update) new_state = TrainingState( optimizer_state=optimizer_state, policy_params=policy_params, key=key, steps=state.steps + 1, ) return new_state, metrics # General learner book-keeping and loggers. self._counter = counter or counting.Counter(prefix='learner') self._logger = logger or loggers.make_default_logger( 'learner', asynchronous=True, serialize_fn=utils.fetch_devicearray, steps_key=self._counter.get_steps_key()) # Split the input batch to `num_sgd_steps_per_step` minibatches in order # to achieve better performance on accelerators. sgd_step = utils.process_multiple_batches(sgd_step, num_sgd_steps_per_step) self._sgd_step = jax.pmap(sgd_step, axis_name=_PMAP_AXIS_NAME) random_key, init_key = jax.random.split(random_key) policy_params = network.init(init_key) optimizer_state = optimizer.init(policy_params) # Create initial state. state = TrainingState( optimizer_state=optimizer_state, policy_params=policy_params, key=random_key, steps=0, ) self._state = utils.replicate_in_all_devices(state) self._timestamp = None self._prefetching_iterator = prefetching_iterator
def __init__(self, network: networks_lib.FeedForwardNetwork, random_key: networks_lib.PRNGKey, loss_fn: losses.Loss, optimizer: optax.GradientTransformation, demonstrations: Iterator[types.Transition], num_sgd_steps_per_step: int, logger: Optional[loggers.Logger] = None, counter: Optional[counting.Counter] = None): def sgd_step( state: TrainingState, transitions: types.Transition, ) -> Tuple[TrainingState, Dict[str, jnp.ndarray]]: loss_and_grad = jax.value_and_grad(loss_fn, argnums=1) # Compute losses and their gradients. key, key_input = jax.random.split(state.key) loss_value, gradients = loss_and_grad(network.apply, state.policy_params, key_input, transitions) policy_update, optimizer_state = optimizer.update(gradients, state.optimizer_state) policy_params = optax.apply_updates(state.policy_params, policy_update) new_state = TrainingState( optimizer_state=optimizer_state, policy_params=policy_params, key=key, steps=state.steps + 1, ) metrics = { 'loss': loss_value, 'gradient_norm': optax.global_norm(gradients) } return new_state, metrics # General learner book-keeping and loggers. self._counter = counter or counting.Counter(prefix='learner') self._logger = logger or loggers.make_default_logger( 'learner', asynchronous=True, serialize_fn=utils.fetch_devicearray) # Iterator on demonstration transitions. self._demonstrations = demonstrations # Split the input batch to `num_sgd_steps_per_step` minibatches in order # to achieve better performance on accelerators. self._sgd_step = jax.jit(utils.process_multiple_batches( sgd_step, num_sgd_steps_per_step)) random_key, init_key = jax.random.split(random_key) policy_params = network.init(init_key) optimizer_state = optimizer.init(policy_params) # Create initial state. self._state = TrainingState( optimizer_state=optimizer_state, policy_params=policy_params, key=random_key, steps=0, ) self._timestamp = None