def flush(self): logger.info("Loss on {} batches".format(len(self.incoming_td_loss))) print_details = "Loss:\n" td_loss = torch.tensor(self.incoming_td_loss) SummaryWriterContext.add_histogram("td_loss", td_loss) td_loss_mean = float(td_loss.mean()) SummaryWriterContext.add_scalar("td_loss/mean", td_loss_mean) self.td_loss.append(td_loss_mean) print_details = print_details + "TD LOSS: {0:.3f}\n".format( td_loss_mean) if len(self.incoming_reward_loss) > 0: reward_loss = torch.tensor(self.incoming_reward_loss) SummaryWriterContext.add_histogram("reward_loss", reward_loss) reward_loss_mean = float(reward_loss.mean()) SummaryWriterContext.add_scalar("reward_loss/mean", reward_loss_mean) self.reward_loss.append(reward_loss_mean) print_details = print_details + "REWARD LOSS: {0:.3f}\n".format( reward_loss_mean) for print_detail in print_details.split("\n"): logger.info(print_detail) self.incoming_td_loss.clear() self.incoming_reward_loss.clear()
def _log_histogram_and_mean(self, log_key, val): try: SummaryWriterContext.add_histogram(log_key, val) SummaryWriterContext.add_scalar(f"{log_key}/mean", val.mean()) except ValueError: logger.warning( f"Cannot create histogram for key: {log_key}; " "this is likely because you have NULL value in your input; " f"value: {val}") raise
def get_log_prob(self, state, squashed_action): """ Action is expected to be squashed with tanh """ loc, scale_log = self._get_loc_and_scale_log(state) # This is not getting exported; we can use it n = Normal(loc, scale_log.exp()) raw_action = self._atanh(squashed_action) log_prob = n.log_prob(raw_action) squash_correction = self._squash_correction(squashed_action) if SummaryWriterContext._global_step % 1000 == 0: SummaryWriterContext.add_histogram("actor/get_log_prob/loc", loc.detach().cpu()) SummaryWriterContext.add_histogram("actor/get_log_prob/scale_log", scale_log.detach().cpu()) SummaryWriterContext.add_histogram("actor/get_log_prob/log_prob", log_prob.detach().cpu()) SummaryWriterContext.add_histogram( "actor/get_log_prob/squash_correction", squash_correction.detach().cpu()) log_prob = torch.sum(log_prob - squash_correction, dim=1).reshape(-1, 1) return log_prob
def forward(self, input): loc, scale_log = self._get_loc_and_scale_log(input.state) r = torch.randn_like(scale_log, device=scale_log.device) action = torch.tanh(loc + r * scale_log.exp()) if not self.training: # ONNX doesn't like reshape either.. return rlt.ActorOutput(action=action) # Since each dim are independent, log-prob is simply sum log_prob = self._log_prob(r, scale_log) squash_correction = self._squash_correction(action) if SummaryWriterContext._global_step % 1000 == 0: SummaryWriterContext.add_histogram("actor/forward/loc", loc.detach().cpu()) SummaryWriterContext.add_histogram("actor/forward/scale_log", scale_log.detach().cpu()) SummaryWriterContext.add_histogram("actor/forward/log_prob", log_prob.detach().cpu()) SummaryWriterContext.add_histogram( "actor/forward/squash_correction", squash_correction.detach().cpu()) log_prob = torch.sum(log_prob - squash_correction, dim=1) return rlt.ActorOutput(action=action, log_prob=log_prob.reshape(-1, 1), action_mean=loc)
def forward(self, input) -> torch.FloatTensor: state_dim = self.layers[0].in_features state = input[:, :state_dim] action = input[:, state_dim:] x = state for i, activation in enumerate(self.activations[:-1]): if self.use_batch_norm: x = self.batch_norm_ops[i](x) activation_func = getattr(F, activation) fc_func = self.layers[i] x = fc_func(x) if activation == "linear" else activation_func(fc_func(x)) value = self.value(x) x = torch.cat((x, action), dim=1) raw_advantage = self.advantage(x) if self.parametric_action: advantage = raw_advantage else: advantage = raw_advantage - raw_advantage.mean(dim=1, keepdim=True) q_value = value + advantage if SummaryWriterContext._global_step % 1000 == 0: SummaryWriterContext.add_histogram( "dueling_network/{}/value".format(self._name), value.detach().cpu() ) SummaryWriterContext.add_scalar( "dueling_network/{}/mean_value".format(self._name), value.detach().mean().cpu(), ) SummaryWriterContext.add_histogram( "dueling_network/{}/q_value".format(self._name), q_value.detach().cpu() ) SummaryWriterContext.add_scalar( "dueling_network/{}/mean_q_value".format(self._name), q_value.detach().mean().cpu(), ) SummaryWriterContext.add_histogram( "dueling_network/{}/raw_advantage".format(self._name), raw_advantage.detach().cpu(), ) SummaryWriterContext.add_scalar( "dueling_network/{}/mean_raw_advantage".format(self._name), raw_advantage.detach().mean().cpu(), ) if not self.parametric_action: for i in range(advantage.shape[1]): a = advantage.detach()[:, i] SummaryWriterContext.add_histogram( "dueling_network/{}/advatage/{}".format(self._name, i), a.cpu() ) SummaryWriterContext.add_scalar( "dueling_network/{}/mean_advatage/{}".format(self._name, i), a.mean().cpu(), ) return q_value
def write_summary(self, actions: List[str]): if actions: for field, log_key in [ ("logged_actions", "actions/logged"), ("model_action_idxs", "actions/model"), ]: val = getattr(self, field) if val is None: continue for i, action in enumerate(actions): SummaryWriterContext.add_scalar( "{}/{}".format(log_key, action), (val == i).sum().item() ) for field, log_key in [ ("td_loss", "td_loss"), ("reward_loss", "reward_loss"), ("logged_propensities", "propensities/logged"), ("logged_rewards", "reward/logged"), ("logged_values", "value/logged"), ("model_values_on_logged_actions", "value/model_logged_action"), ]: val = getattr(self, field) if val is None: continue assert len(val.shape) == 1 or ( len(val.shape) == 2 and val.shape[1] == 1 ), "Unexpected shape for {}: {}".format(field, val.shape) SummaryWriterContext.add_histogram(log_key, val) SummaryWriterContext.add_scalar("{}/mean".format(log_key), val.mean()) for field, log_key in [ ("model_propensities", "propensities/model"), ("model_rewards", "reward/model"), ("model_values", "value/model"), ]: val = getattr(self, field) if val is None: continue if ( len(val.shape) == 1 or (len(val.shape) == 2 and val.shape[1] == 1) ) and not actions: SummaryWriterContext.add_histogram(log_key, val) SummaryWriterContext.add_scalar("{}/mean".format(log_key), val.mean()) elif len(val.shape) == 2 and val.shape[1] == len(actions): for i, action in enumerate(actions): SummaryWriterContext.add_histogram( "{}/{}".format(log_key, action), val[:, i] ) SummaryWriterContext.add_scalar( "{}/{}/mean".format(log_key, action), val[:, i].mean() ) else: raise ValueError( "Unexpected shape for {}: {}; actions: {}".format( field, val.shape, actions ) )
def _sample_action(self, loc: torch.Tensor, scale_log: torch.Tensor): r = torch.randn_like(scale_log, device=scale_log.device) action = torch.tanh(loc + r * scale_log.exp()) # Since each dim are independent, log-prob is simply sum log_prob = self.actor_network._log_prob(r, scale_log) squash_correction = self.actor_network._squash_correction(action) if SummaryWriterContext._global_step % 1000 == 0: SummaryWriterContext.add_histogram("actor/forward/loc", loc.detach().cpu()) SummaryWriterContext.add_histogram("actor/forward/scale_log", scale_log.detach().cpu()) SummaryWriterContext.add_histogram("actor/forward/log_prob", log_prob.detach().cpu()) SummaryWriterContext.add_histogram( "actor/forward/squash_correction", squash_correction.detach().cpu()) log_prob = torch.sum(log_prob - squash_correction, dim=1) return action, log_prob.reshape(-1, 1)
def _log_prob(self, loc: torch.Tensor, scale_log: torch.Tensor, squashed_action: torch.Tensor): # This is not getting exported; we can use it n = torch.distributions.Normal(loc, scale_log.exp()) raw_action = self.actor_network._atanh(squashed_action) log_prob = n.log_prob(raw_action) squash_correction = self.actor_network._squash_correction( squashed_action) if SummaryWriterContext._global_step % 1000 == 0: SummaryWriterContext.add_histogram("actor/get_log_prob/loc", loc.detach().cpu()) SummaryWriterContext.add_histogram("actor/get_log_prob/scale_log", scale_log.detach().cpu()) SummaryWriterContext.add_histogram("actor/get_log_prob/log_prob", log_prob.detach().cpu()) SummaryWriterContext.add_histogram( "actor/get_log_prob/squash_correction", squash_correction.detach().cpu()) log_prob = torch.sum(log_prob - squash_correction, dim=1).reshape(-1, 1) return log_prob
def train(self, training_batch) -> None: """ IMPORTANT: the input action here is assumed to be preprocessed to match the range of the output of the actor. """ if hasattr(training_batch, "as_parametric_sarsa_training_batch"): training_batch = training_batch.as_parametric_sarsa_training_batch( ) learning_input = training_batch.training_input self.minibatch += 1 state = learning_input.state action = learning_input.action next_state = learning_input.next_state reward = learning_input.reward not_done_mask = learning_input.not_terminal action = self._maybe_scale_action_in_train(action) # Compute current value estimates current_state_action = rlt.StateAction(state=state, action=action) q1_value = self.q1_network(current_state_action).q_value if self.q2_network: q2_value = self.q2_network(current_state_action).q_value actor_action = self.actor_network(rlt.StateInput(state=state)).action # Generate target = r + y * min (Q1(s',pi(s')), Q2(s',pi(s'))) with torch.no_grad(): next_actor = self.actor_network_target( rlt.StateInput(state=next_state)).action next_actor += (torch.randn_like(next_actor) * self.target_policy_smoothing).clamp( -self.noise_clip, self.noise_clip) next_actor = torch.max( torch.min(next_actor, self.max_action_range_tensor_training), self.min_action_range_tensor_training, ) next_state_actor = rlt.StateAction( state=next_state, action=rlt.FeatureVector(float_features=next_actor)) next_state_value = self.q1_network_target(next_state_actor).q_value if self.q2_network is not None: next_state_value = torch.min( next_state_value, self.q2_network_target(next_state_actor).q_value) target_q_value = ( reward + self.gamma * next_state_value * not_done_mask.float()) # Optimize Q1 and Q2 q1_loss = F.mse_loss(q1_value, target_q_value) q1_loss.backward() self._maybe_run_optimizer(self.q1_network_optimizer, self.minibatches_per_step) if self.q2_network: q2_loss = F.mse_loss(q2_value, target_q_value) q2_loss.backward() self._maybe_run_optimizer(self.q2_network_optimizer, self.minibatches_per_step) # Only update actor and target networks after a fixed number of Q updates if self.minibatch % self.delayed_policy_update == 0: actor_loss = -self.q1_network( rlt.StateAction( state=state, action=rlt.FeatureVector( float_features=actor_action))).q_value.mean() actor_loss.backward() self._maybe_run_optimizer(self.actor_network_optimizer, self.minibatches_per_step) # Use the soft update rule to update the target networks self._maybe_soft_update( self.q1_network, self.q1_network_target, self.tau, self.minibatches_per_step, ) self._maybe_soft_update( self.actor_network, self.actor_network_target, self.tau, self.minibatches_per_step, ) if self.q2_network is not None: self._maybe_soft_update( self.q2_network, self.q2_network_target, self.tau, self.minibatches_per_step, ) # Logging at the end to schedule all the cuda operations first if (self.tensorboard_logging_freq is not None and self.minibatch % self.tensorboard_logging_freq == 0): SummaryWriterContext.add_histogram("q1/logged_state_value", q1_value) if self.q2_network: SummaryWriterContext.add_histogram("q2/logged_state_value", q2_value) SummaryWriterContext.add_histogram("q_network/next_state_value", next_state_value) SummaryWriterContext.add_histogram("q_network/target_q_value", target_q_value) SummaryWriterContext.add_histogram("actor/loss", actor_loss) self.loss_reporter.report( td_loss=float(q1_loss), reward_loss=None, logged_rewards=reward, model_values_on_logged_actions=q1_value, )
def train(self, training_batch, evaluator=None) -> None: """ IMPORTANT: the input action here is assumed to be preprocessed to match the range of the output of the actor. """ if hasattr(training_batch, "as_parametric_sarsa_training_batch"): training_batch = training_batch.as_parametric_sarsa_training_batch( ) learning_input = training_batch.training_input self.minibatch += 1 state = learning_input.state action = learning_input.action reward = learning_input.reward discount = torch.full_like(reward, self.gamma) not_done_mask = learning_input.not_terminal if self._should_scale_action_in_train(): action = rlt.FeatureVector( rescale_torch_tensor( action.float_features, new_min=self.min_action_range_tensor_training, new_max=self.max_action_range_tensor_training, prev_min=self.min_action_range_tensor_serving, prev_max=self.max_action_range_tensor_serving, )) current_state_action = rlt.StateAction(state=state, action=action) q1_value = self.q1_network(current_state_action).q_value min_q_value = q1_value if self.q2_network: q2_value = self.q2_network(current_state_action).q_value min_q_value = torch.min(q1_value, q2_value) # Use the minimum as target, ensure no gradient going through min_q_value = min_q_value.detach() # # First, optimize value network; minimizing MSE between # V(s) & Q(s, a) - log(pi(a|s)) # state_value = self.value_network(state.float_features) # .q_value if self.logged_action_uniform_prior: log_prob_a = torch.zeros_like(min_q_value) target_value = min_q_value else: with torch.no_grad(): log_prob_a = self.actor_network.get_log_prob( state, action.float_features) log_prob_a = log_prob_a.clamp(-20.0, 20.0) target_value = min_q_value - self.entropy_temperature * log_prob_a value_loss = F.mse_loss(state_value, target_value) self.value_network_optimizer.zero_grad() value_loss.backward() self.value_network_optimizer.step() # # Second, optimize Q networks; minimizing MSE between # Q(s, a) & r + discount * V'(next_s) # with torch.no_grad(): next_state_value = (self.value_network_target( learning_input.next_state.float_features) * not_done_mask) if self.minibatch < self.reward_burnin: target_q_value = reward else: target_q_value = reward + discount * next_state_value q1_loss = F.mse_loss(q1_value, target_q_value) self.q1_network_optimizer.zero_grad() q1_loss.backward() self.q1_network_optimizer.step() if self.q2_network: q2_loss = F.mse_loss(q2_value, target_q_value) self.q2_network_optimizer.zero_grad() q2_loss.backward() self.q2_network_optimizer.step() # # Lastly, optimize the actor; minimizing KL-divergence between action propensity # & softmax of value. Due to reparameterization trick, it ends up being # log_prob(actor_action) - Q(s, actor_action) # actor_output = self.actor_network(rlt.StateInput(state=state)) state_actor_action = rlt.StateAction( state=state, action=rlt.FeatureVector(float_features=actor_output.action)) q1_actor_value = self.q1_network(state_actor_action).q_value min_q_actor_value = q1_actor_value if self.q2_network: q2_actor_value = self.q2_network(state_actor_action).q_value min_q_actor_value = torch.min(q1_actor_value, q2_actor_value) actor_loss = (self.entropy_temperature * actor_output.log_prob - min_q_actor_value) # Do this in 2 steps so we can log histogram of actor loss actor_loss_mean = actor_loss.mean() self.actor_network_optimizer.zero_grad() actor_loss_mean.backward() self.actor_network_optimizer.step() if self.minibatch < self.reward_burnin: # Reward burnin: force target network self._soft_update(self.value_network, self.value_network_target, 1.0) else: # Use the soft update rule to update both target networks self._soft_update(self.value_network, self.value_network_target, self.tau) # Logging at the end to schedule all the cuda operations first if (self.tensorboard_logging_freq is not None and self.minibatch % self.tensorboard_logging_freq == 0): SummaryWriterContext.add_histogram("q1/logged_state_value", q1_value) if self.q2_network: SummaryWriterContext.add_histogram("q2/logged_state_value", q2_value) SummaryWriterContext.add_histogram("log_prob_a", log_prob_a) SummaryWriterContext.add_histogram("value_network/target", target_value) SummaryWriterContext.add_histogram("q_network/next_state_value", next_state_value) SummaryWriterContext.add_histogram("q_network/target_q_value", target_q_value) SummaryWriterContext.add_histogram("actor/min_q_actor_value", min_q_actor_value) SummaryWriterContext.add_histogram("actor/action_log_prob", actor_output.log_prob) SummaryWriterContext.add_histogram("actor/loss", actor_loss) if evaluator is not None: cpe_stats = BatchStatsForCPE( td_loss=q1_loss.detach().cpu().numpy(), logged_rewards=reward.detach().cpu().numpy(), model_values_on_logged_actions=q1_value.detach().cpu().numpy(), model_propensities=actor_output.log_prob.exp().detach().cpu(). numpy(), model_values=min_q_actor_value.detach().cpu().numpy(), ) evaluator.report(cpe_stats)
def dist(self, input: rlt.PreprocessedState): state = input.state.float_features x = state for i, activation in enumerate(self.activations[:-1]): if self.use_batch_norm: x = self.batch_norm_ops[i](x) x = self.layers[i](x) if activation == "linear": continue elif activation == "tanh": activation_func = torch.tanh else: activation_func = getattr(F, activation) x = activation_func(x) value = self.value(x).unsqueeze(dim=1) raw_advantage = self.advantage(x).reshape(-1, self.num_actions, self.num_atoms) advantage = raw_advantage - raw_advantage.mean(dim=1, keepdim=True) q_value = value + advantage if SummaryWriterContext._global_step % 1000 == 0: SummaryWriterContext.add_histogram( "dueling_network/{}/value".format(self._name), value.detach().mean(dim=2).cpu(), ) SummaryWriterContext.add_scalar( "dueling_network/{}/mean_value".format(self._name), value.detach().mean().cpu(), ) SummaryWriterContext.add_histogram( "dueling_network/{}/q_value".format(self._name), q_value.detach().mean(dim=2).cpu(), ) SummaryWriterContext.add_scalar( "dueling_network/{}/mean_q_value".format(self._name), q_value.detach().mean().cpu(), ) SummaryWriterContext.add_histogram( "dueling_network/{}/raw_advantage".format(self._name), raw_advantage.detach().mean(dim=2).cpu(), ) SummaryWriterContext.add_scalar( "dueling_network/{}/mean_raw_advantage".format(self._name), raw_advantage.detach().mean().cpu(), ) for i in range(advantage.shape[1]): a = advantage.detach()[:, i, :].mean(dim=1) SummaryWriterContext.add_histogram( "dueling_network/{}/advantage/{}".format(self._name, i), a.cpu()) SummaryWriterContext.add_scalar( "dueling_network/{}/mean_advantage/{}".format( self._name, i), a.mean().cpu(), ) return q_value
def forward(self, input) -> torch.FloatTensor: output_tensor = False if isinstance(input, torch.Tensor): # Maintaining backward compatibility for a bit state_dim = self.layers[0].in_features state = input[:, :state_dim] action = input[:, state_dim:] output_tensor = True elif self.parametric_action: state = input.state.float_features action = input.action.float_features else: state = input.state.float_features action = None x = state for i, activation in enumerate(self.activations[:-1]): if self.use_batch_norm: x = self.batch_norm_ops[i](x) activation_func = getattr(F, activation) fc_func = self.layers[i] x = fc_func(x) if activation == "linear" else activation_func(fc_func(x)) value = self.value(x) if action is not None: x = torch.cat((x, action), dim=1) raw_advantage = self.advantage(x) if self.parametric_action: advantage = raw_advantage else: advantage = raw_advantage - raw_advantage.mean(dim=1, keepdim=True) q_value = value + advantage if SummaryWriterContext._global_step % 1000 == 0: SummaryWriterContext.add_histogram( "dueling_network/{}/value".format(self._name), value.detach().cpu() ) SummaryWriterContext.add_scalar( "dueling_network/{}/mean_value".format(self._name), value.detach().mean().cpu(), ) SummaryWriterContext.add_histogram( "dueling_network/{}/q_value".format(self._name), q_value.detach().cpu() ) SummaryWriterContext.add_scalar( "dueling_network/{}/mean_q_value".format(self._name), q_value.detach().mean().cpu(), ) SummaryWriterContext.add_histogram( "dueling_network/{}/raw_advantage".format(self._name), raw_advantage.detach().cpu(), ) SummaryWriterContext.add_scalar( "dueling_network/{}/mean_raw_advantage".format(self._name), raw_advantage.detach().mean().cpu(), ) if not self.parametric_action: for i in range(advantage.shape[1]): a = advantage.detach()[:, i] SummaryWriterContext.add_histogram( "dueling_network/{}/advatage/{}".format(self._name, i), a.cpu() ) SummaryWriterContext.add_scalar( "dueling_network/{}/mean_advatage/{}".format(self._name, i), a.mean().cpu(), ) if output_tensor: return q_value elif self.parametric_action: return rlt.SingleQValue(q_value=q_value) else: return rlt.AllActionQValues(q_values=q_value)
def train(self, training_batch) -> None: """ IMPORTANT: the input action here is assumed to be preprocessed to match the range of the output of the actor. """ if hasattr(training_batch, "as_parametric_sarsa_training_batch"): training_batch = training_batch.as_parametric_sarsa_training_batch() learning_input = training_batch.training_input self.minibatch += 1 state = learning_input.state action = learning_input.action reward = learning_input.reward discount = torch.full_like(reward, self.gamma) not_done_mask = learning_input.not_terminal if self._should_scale_action_in_train(): action = rlt.FeatureVector( rescale_torch_tensor( action.float_features, new_min=self.min_action_range_tensor_training, new_max=self.max_action_range_tensor_training, prev_min=self.min_action_range_tensor_serving, prev_max=self.max_action_range_tensor_serving, ) ) current_state_action = rlt.StateAction(state=state, action=action) q1_value = self.q1_network(current_state_action).q_value min_q_value = q1_value if self.q2_network: q2_value = self.q2_network(current_state_action).q_value min_q_value = torch.min(q1_value, q2_value) # Use the minimum as target, ensure no gradient going through min_q_value = min_q_value.detach() # # First, optimize value network; minimizing MSE between # V(s) & Q(s, a) - log(pi(a|s)) # state_value = self.value_network(state.float_features) # .q_value if self.logged_action_uniform_prior: log_prob_a = torch.zeros_like(min_q_value) target_value = min_q_value else: with torch.no_grad(): log_prob_a = self.actor_network.get_log_prob( state, action.float_features ) log_prob_a = log_prob_a.clamp(-20.0, 20.0) target_value = min_q_value - self.entropy_temperature * log_prob_a value_loss = F.mse_loss(state_value, target_value) self.value_network_optimizer.zero_grad() value_loss.backward() self.value_network_optimizer.step() # # Second, optimize Q networks; minimizing MSE between # Q(s, a) & r + discount * V'(next_s) # with torch.no_grad(): next_state_value = ( self.value_network_target(learning_input.next_state.float_features) * not_done_mask.float() ) if self.minibatch < self.reward_burnin: target_q_value = reward else: target_q_value = reward + discount * next_state_value q1_loss = F.mse_loss(q1_value, target_q_value) self.q1_network_optimizer.zero_grad() q1_loss.backward() self.q1_network_optimizer.step() if self.q2_network: q2_loss = F.mse_loss(q2_value, target_q_value) self.q2_network_optimizer.zero_grad() q2_loss.backward() self.q2_network_optimizer.step() # # Lastly, optimize the actor; minimizing KL-divergence between action propensity # & softmax of value. Due to reparameterization trick, it ends up being # log_prob(actor_action) - Q(s, actor_action) # actor_output = self.actor_network(rlt.StateInput(state=state)) state_actor_action = rlt.StateAction( state=state, action=rlt.FeatureVector(float_features=actor_output.action) ) q1_actor_value = self.q1_network(state_actor_action).q_value min_q_actor_value = q1_actor_value if self.q2_network: q2_actor_value = self.q2_network(state_actor_action).q_value min_q_actor_value = torch.min(q1_actor_value, q2_actor_value) actor_loss = ( self.entropy_temperature * actor_output.log_prob - min_q_actor_value ) # Do this in 2 steps so we can log histogram of actor loss actor_loss_mean = actor_loss.mean() self.actor_network_optimizer.zero_grad() actor_loss_mean.backward() self.actor_network_optimizer.step() if self.minibatch < self.reward_burnin: # Reward burnin: force target network self._soft_update(self.value_network, self.value_network_target, 1.0) else: # Use the soft update rule to update both target networks self._soft_update(self.value_network, self.value_network_target, self.tau) # Logging at the end to schedule all the cuda operations first if ( self.tensorboard_logging_freq is not None and self.minibatch % self.tensorboard_logging_freq == 0 ): SummaryWriterContext.add_histogram("q1/logged_state_value", q1_value) if self.q2_network: SummaryWriterContext.add_histogram("q2/logged_state_value", q2_value) SummaryWriterContext.add_histogram("log_prob_a", log_prob_a) SummaryWriterContext.add_histogram("value_network/target", target_value) SummaryWriterContext.add_histogram( "q_network/next_state_value", next_state_value ) SummaryWriterContext.add_histogram( "q_network/target_q_value", target_q_value ) SummaryWriterContext.add_histogram( "actor/min_q_actor_value", min_q_actor_value ) SummaryWriterContext.add_histogram( "actor/action_log_prob", actor_output.log_prob ) SummaryWriterContext.add_histogram("actor/loss", actor_loss) self.loss_reporter.report( td_loss=float(q1_loss), reward_loss=None, logged_rewards=reward, model_values_on_logged_actions=q1_value, model_propensities=actor_output.log_prob.exp(), model_values=min_q_actor_value, )
def evaluate_batch(self): merged_inputs = [] for batch in self.all_batches: if len(batch) > 0: merged_inputs.append(np.vstack(batch)) else: merged_inputs.append(None) ( td_loss, logged_actions, logged_propensities, logged_rewards, logged_values, model_propensities, model_rewards, model_values, model_values_on_logged_actions, model_action_idxs, ) = merged_inputs logger.info("Evaluating on {} batches".format(len( self.td_loss_batches))) print_details = "Evaluator:\n" if td_loss is not None: SummaryWriterContext.add_histogram("td_loss", td_loss) td_loss_mean = float(np.mean(td_loss)) SummaryWriterContext.add_scalar("td_loss/mean", td_loss_mean) self.td_loss.append(td_loss_mean) print_details = print_details + "TD LOSS: {0:.3f}\n".format( td_loss_mean) if logged_rewards is not None: SummaryWriterContext.add_histogram("reward/logged", logged_rewards) SummaryWriterContext.add_scalar("reward/logged/mean", logged_rewards.mean()) if model_rewards is not None: SummaryWriterContext.add_histogram("reward/model", model_rewards) SummaryWriterContext.add_scalar("reward/model/mean", model_rewards.mean()) if logged_values is not None: SummaryWriterContext.add_histogram("value/logged", logged_values) SummaryWriterContext.add_scalar("value/logged/mean", logged_values.mean()) if model_values is not None: SummaryWriterContext.add_histogram("value/model", model_values) SummaryWriterContext.add_scalar("value/model/mean", model_values.mean()) if model_values_on_logged_actions is not None: SummaryWriterContext.add_histogram("value/model_logged_action", model_values_on_logged_actions) # TODO: log summary of logged propensities if model_propensities is not None and self.action_names: if len(model_propensities.shape) == 1: SummaryWriterContext.add_histogram("propensities/model", model_propensities) SummaryWriterContext.add_scalar("propensities/model/mean", model_propensities.mean()) if len(model_propensities.shape) == 2: for i, action_name in enumerate(self.action_names): SummaryWriterContext.add_histogram( "propensities/model/{}".format(action_name), model_propensities[:, i], ) SummaryWriterContext.add_scalar( "propensities/model/{}/mean".format(action_name), model_propensities[:, i].mean(), ) if logged_actions is not None and model_action_idxs is not None: logged_action_distr, logged_action_counts = self._get_batch_logged_actions( [logged_actions]) model_action_distr, model_action_counts = self._get_batch_model_actions( [model_action_idxs]) print_details += "The distribution of logged actions : {}\n".format( logged_action_counts) print_details += "The distribution of model actions : {}\n".format( model_action_counts) for action, count in logged_action_counts.items(): self.logged_action_counts[action] += count for action, count in model_action_counts.items(): self.model_action_counts[action].append(count) self.model_action_counts_cumulative[action] += count for action, val in model_action_distr.items(): self.model_action_distr[action].append(val) # Log to tensorboard for action_name, count in logged_action_counts.items(): SummaryWriterContext.add_scalar( "actions/logged/{}".format(action_name), count) for action_name, count in model_action_counts.items(): SummaryWriterContext.add_scalar( "actions/model/{}".format(action_name), count) print_details += "Batch Evaluator Finished" for print_detail in print_details.split("\n"): logger.info(print_detail)
def forward(self, input) -> Union[NamedTuple, torch.FloatTensor]: # type: ignore output_tensor = False if self.parametric_action: state = input.state.float_features action = input.action.float_features else: state = input.state.float_features action = None x = state for i, activation in enumerate(self.activations[:-1]): if self.use_batch_norm: x = self.batch_norm_ops[i](x) x = self.layers[i](x) if activation == "linear": continue elif activation == "tanh": activation_func = torch.tanh else: activation_func = getattr(F, activation) x = activation_func(x) value = self.value(x) if action is not None: x = torch.cat((x, action), dim=1) raw_advantage = self.advantage(x) if self.parametric_action: advantage = raw_advantage else: advantage = raw_advantage - raw_advantage.mean(dim=1, keepdim=True) q_value = value + advantage if SummaryWriterContext._global_step % 1000 == 0: SummaryWriterContext.add_histogram( "dueling_network/{}/value".format(self._name), value.detach().cpu()) SummaryWriterContext.add_scalar( "dueling_network/{}/mean_value".format(self._name), value.detach().mean().cpu(), ) SummaryWriterContext.add_histogram( "dueling_network/{}/q_value".format(self._name), q_value.detach().cpu()) SummaryWriterContext.add_scalar( "dueling_network/{}/mean_q_value".format(self._name), q_value.detach().mean().cpu(), ) SummaryWriterContext.add_histogram( "dueling_network/{}/raw_advantage".format(self._name), raw_advantage.detach().cpu(), ) SummaryWriterContext.add_scalar( "dueling_network/{}/mean_raw_advantage".format(self._name), raw_advantage.detach().mean().cpu(), ) if not self.parametric_action: for i in range(advantage.shape[1]): a = advantage.detach()[:, i] SummaryWriterContext.add_histogram( "dueling_network/{}/advantage/{}".format( self._name, i), a.cpu()) SummaryWriterContext.add_scalar( "dueling_network/{}/mean_advantage/{}".format( self._name, i), a.mean().cpu(), ) if output_tensor: return q_value # type: ignore elif self.parametric_action: return rlt.SingleQValue(q_value=q_value) # type: ignore else: return rlt.AllActionQValues(q_values=q_value) # type: ignore
def train(self, training_batch) -> None: """ IMPORTANT: the input action here is assumed to be preprocessed to match the range of the output of the actor. """ if hasattr(training_batch, "as_parametric_sarsa_training_batch"): training_batch = training_batch.as_parametric_sarsa_training_batch( ) learning_input = training_batch.training_input self.minibatch += 1 state = learning_input.state action = learning_input.action reward = learning_input.reward discount = torch.full_like(reward, self.gamma) not_done_mask = learning_input.not_terminal if self._should_scale_action_in_train(): action = rlt.FeatureVector( rescale_torch_tensor( action.float_features, new_min=self.min_action_range_tensor_training, new_max=self.max_action_range_tensor_training, prev_min=self.min_action_range_tensor_serving, prev_max=self.max_action_range_tensor_serving, )) with torch.enable_grad(): # # First, optimize Q networks; minimizing MSE between # Q(s, a) & r + discount * V'(next_s) # current_state_action = rlt.StateAction(state=state, action=action) q1_value = self.q1_network(current_state_action).q_value if self.q2_network: q2_value = self.q2_network(current_state_action).q_value actor_output = self.actor_network(rlt.StateInput(state=state)) # Optimize Alpha if self.alpha_optimizer is not None: alpha_loss = -(self.log_alpha * (actor_output.log_prob + self.target_entropy).detach()).mean() self.alpha_optimizer.zero_grad() alpha_loss.backward() self.alpha_optimizer.step() self.entropy_temperature = self.log_alpha.exp() with torch.no_grad(): if self.value_network is not None: next_state_value = self.value_network_target( learning_input.next_state.float_features) else: next_state_actor_output = self.actor_network( rlt.StateInput(state=learning_input.next_state)) next_state_actor_action = rlt.StateAction( state=learning_input.next_state, action=rlt.FeatureVector( float_features=next_state_actor_output.action), ) next_state_value = self.q1_network_target( next_state_actor_action).q_value if self.q2_network is not None: target_q2_value = self.q2_network_target( next_state_actor_action).q_value next_state_value = torch.min(next_state_value, target_q2_value) log_prob_a = self.actor_network.get_log_prob( learning_input.next_state, next_state_actor_output.action) log_prob_a = log_prob_a.clamp(-20.0, 20.0) next_state_value -= self.entropy_temperature * log_prob_a target_q_value = ( reward + discount * next_state_value * not_done_mask.float()) q1_loss = F.mse_loss(q1_value, target_q_value) q1_loss.backward() self._maybe_run_optimizer(self.q1_network_optimizer, self.minibatches_per_step) if self.q2_network: q2_loss = F.mse_loss(q2_value, target_q_value) q2_loss.backward() self._maybe_run_optimizer(self.q2_network_optimizer, self.minibatches_per_step) # # Second, optimize the actor; minimizing KL-divergence between action propensity # & softmax of value. Due to reparameterization trick, it ends up being # log_prob(actor_action) - Q(s, actor_action) # state_actor_action = rlt.StateAction( state=state, action=rlt.FeatureVector(float_features=actor_output.action), ) q1_actor_value = self.q1_network(state_actor_action).q_value min_q_actor_value = q1_actor_value if self.q2_network: q2_actor_value = self.q2_network(state_actor_action).q_value min_q_actor_value = torch.min(q1_actor_value, q2_actor_value) actor_loss = (self.entropy_temperature * actor_output.log_prob - min_q_actor_value) # Do this in 2 steps so we can log histogram of actor loss actor_loss_mean = actor_loss.mean() actor_loss_mean.backward() self._maybe_run_optimizer(self.actor_network_optimizer, self.minibatches_per_step) # # Lastly, if applicable, optimize value network; minimizing MSE between # V(s) & E_a~pi(s) [ Q(s,a) - log(pi(a|s)) ] # if self.value_network is not None: state_value = self.value_network(state.float_features) if self.logged_action_uniform_prior: log_prob_a = torch.zeros_like(min_q_actor_value) target_value = min_q_actor_value else: with torch.no_grad(): log_prob_a = actor_output.log_prob log_prob_a = log_prob_a.clamp(-20.0, 20.0) target_value = (min_q_actor_value - self.entropy_temperature * log_prob_a) value_loss = F.mse_loss(state_value, target_value.detach()) value_loss.backward() self._maybe_run_optimizer(self.value_network_optimizer, self.minibatches_per_step) # Use the soft update rule to update the target networks if self.value_network is not None: self._maybe_soft_update( self.value_network, self.value_network_target, self.tau, self.minibatches_per_step, ) else: self._maybe_soft_update( self.q1_network, self.q1_network_target, self.tau, self.minibatches_per_step, ) if self.q2_network is not None: self._maybe_soft_update( self.q2_network, self.q2_network_target, self.tau, self.minibatches_per_step, ) # Logging at the end to schedule all the cuda operations first if (self.tensorboard_logging_freq is not None and self.minibatch % self.tensorboard_logging_freq == 0): SummaryWriterContext.add_histogram("q1/logged_state_value", q1_value) if self.q2_network: SummaryWriterContext.add_histogram("q2/logged_state_value", q2_value) SummaryWriterContext.add_histogram("log_prob_a", log_prob_a) if self.value_network: SummaryWriterContext.add_histogram("value_network/target", target_value) SummaryWriterContext.add_histogram("q_network/next_state_value", next_state_value) SummaryWriterContext.add_histogram("q_network/target_q_value", target_q_value) SummaryWriterContext.add_histogram("actor/min_q_actor_value", min_q_actor_value) SummaryWriterContext.add_histogram("actor/action_log_prob", actor_output.log_prob) SummaryWriterContext.add_histogram("actor/loss", actor_loss) self.loss_reporter.report( td_loss=float(q1_loss), reward_loss=None, logged_rewards=reward, model_values_on_logged_actions=q1_value, model_propensities=actor_output.log_prob.exp(), model_values=min_q_actor_value, )