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 test_writing(self): with TemporaryDirectory() as tmp_dir: writer = SummaryWriter(tmp_dir) writer.add_scalar = MagicMock() with summary_writer_context(writer): SummaryWriterContext.add_scalar("test", torch.ones(1)) writer.add_scalar.assert_called_once_with("test", torch.ones(1))
def log_to_tensorboard(self, metric_name: str) -> None: def none_to_zero(x: Optional[float]) -> float: if x is None or math.isnan(x): return 0.0 return x for name, value in [ ( "CPE/{}/Direct Method Reward".format(metric_name), self.direct_method.normalized, ), ( "CPE/{}/IPS Reward".format(metric_name), self.inverse_propensity.normalized, ), ( "CPE/{}/Doubly Robust Reward".format(metric_name), self.doubly_robust.normalized, ), ( "CPE/{}/Sequential Doubly Robust".format(metric_name), self.sequential_doubly_robust.normalized, ), ( "CPE/{}/Weighted Sequential Doubly Robust".format(metric_name), self.weighted_doubly_robust.normalized, ), ("CPE/{}/MAGIC".format(metric_name), self.magic.normalized), ]: SummaryWriterContext.add_scalar(name, none_to_zero(value))
def test_swallowing_exception(self): with TemporaryDirectory() as tmp_dir: writer = SummaryWriter(tmp_dir) writer.add_scalar = MagicMock(side_effect=NotImplementedError("test")) writer.exceptions_to_ignore = (NotImplementedError, KeyError) with summary_writer_context(writer): SummaryWriterContext.add_scalar("test", torch.ones(1))
def test_swallowing_exception(self): with TemporaryDirectory() as tmp_dir: writer = SummaryWriter(tmp_dir) writer.add_scalar = MagicMock(side_effect=NotImplementedError("test")) writer.exceptions_to_ignore = (NotImplementedError, KeyError) with summary_writer_context(writer): SummaryWriterContext.add_scalar("test", torch.ones(1))
def test_not_swallowing_exception(self): with TemporaryDirectory() as tmp_dir: writer = SummaryWriter(tmp_dir) writer.add_scalar = MagicMock(side_effect=NotImplementedError("test")) with self.assertRaisesRegexp( NotImplementedError, "test" ), summary_writer_context(writer): SummaryWriterContext.add_scalar("test", torch.ones(1))
def test_not_swallowing_exception(self): with TemporaryDirectory() as tmp_dir: writer = SummaryWriter(tmp_dir) writer.add_scalar = MagicMock(side_effect=NotImplementedError("test")) with self.assertRaisesRegex( NotImplementedError, "test" ), summary_writer_context(writer): SummaryWriterContext.add_scalar("test", torch.ones(1))
def test_writing(self): with TemporaryDirectory() as tmp_dir: writer = SummaryWriter(tmp_dir) writer.add_scalar = MagicMock() with summary_writer_context(writer): SummaryWriterContext.add_scalar("test", torch.ones(1)) writer.add_scalar.assert_called_once_with( "test", torch.ones(1), global_step=0 )
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 log_to_tensorboard(self, epoch: int) -> None: def none_to_zero(x: Optional[float]) -> float: if x is None or math.isnan(x): return 0.0 return x for name, value in [ ("Training/td_loss", self.get_recent_td_loss()), ("Training/reward_loss", self.get_recent_reward_loss()), ("Training/imitator_loss", self.get_recent_imitator_loss()), ]: SummaryWriterContext.add_scalar(name, none_to_zero(value), epoch)
def test_writing_stack(self): with TemporaryDirectory() as tmp_dir1, TemporaryDirectory( ) as tmp_dir2: writer1 = SummaryWriter(tmp_dir1) writer1.add_scalar = MagicMock() writer2 = SummaryWriter(tmp_dir2) writer2.add_scalar = MagicMock() with summary_writer_context(writer1): with summary_writer_context(writer2): SummaryWriterContext.add_scalar("test2", torch.ones(1)) SummaryWriterContext.add_scalar("test1", torch.zeros(1)) writer1.add_scalar.assert_called_once_with("test1", torch.zeros(1)) writer2.add_scalar.assert_called_once_with("test2", torch.ones(1))
def test_global_step(self): with TemporaryDirectory() as tmp_dir: writer = SummaryWriter(tmp_dir) writer.add_scalar = MagicMock() with summary_writer_context(writer): SummaryWriterContext.add_scalar("test", torch.ones(1)) SummaryWriterContext.increase_global_step() SummaryWriterContext.add_scalar("test", torch.zeros(1)) writer.add_scalar.assert_has_calls([ call("test", torch.ones(1), global_step=0), call("test", torch.zeros(1), global_step=1), ]) self.assertEqual(2, len(writer.add_scalar.mock_calls))
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"), ("imitator_loss", "imitator_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) self._log_histogram_and_mean(log_key, val) 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: self._log_histogram_and_mean(log_key, val) elif len(val.shape) == 2 and val.shape[1] == len(actions): for i, action in enumerate(actions): self._log_histogram_and_mean(f"{log_key}/{action}", val[:, i]) else: raise ValueError( "Unexpected shape for {}: {}; actions: {}".format( field, val.shape, actions))
def test_global_step(self): with TemporaryDirectory() as tmp_dir: writer = SummaryWriter(tmp_dir) writer.add_scalar = MagicMock() with summary_writer_context(writer): SummaryWriterContext.add_scalar("test", torch.ones(1)) SummaryWriterContext.increase_global_step() SummaryWriterContext.add_scalar("test", torch.zeros(1)) writer.add_scalar.assert_has_calls( [ call("test", torch.ones(1), global_step=0), call("test", torch.zeros(1), global_step=1), ] ) self.assertEqual(2, len(writer.add_scalar.mock_calls))
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 test_writing_stack(self): with TemporaryDirectory() as tmp_dir1, TemporaryDirectory() as tmp_dir2: writer1 = SummaryWriter(tmp_dir1) writer1.add_scalar = MagicMock() writer2 = SummaryWriter(tmp_dir2) writer2.add_scalar = MagicMock() with summary_writer_context(writer1): with summary_writer_context(writer2): SummaryWriterContext.add_scalar("test2", torch.ones(1)) SummaryWriterContext.add_scalar("test1", torch.zeros(1)) writer1.add_scalar.assert_called_once_with( "test1", torch.zeros(1), global_step=0 ) writer2.add_scalar.assert_called_once_with( "test2", torch.ones(1), global_step=0 )
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 test_noop(self): self.assertIsNone( SummaryWriterContext.add_scalar("test", torch.ones(1)))
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 test_noop(self): self.assertIsNone(SummaryWriterContext.add_scalar("test", torch.ones(1)))
def aggregate(self, values): for i, action in enumerate(self.actions): SummaryWriterContext.add_scalar( f"{self.log_key}/{action}", (values == i).sum().item() )
def test_with_none(self): with summary_writer_context(None): self.assertIsNone(SummaryWriterContext.add_scalar("test", torch.ones(1)))
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) -> 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 test_with_none(self): with summary_writer_context(None): self.assertIsNone( SummaryWriterContext.add_scalar("test", torch.ones(1)))
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_policy_network_training_batch"): training_batch = training_batch.as_policy_network_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 = action._replace(float_features=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, )) # We need to zero out grad here because gradient from actor update # should not be used in Q-network update self.actor_network_optimizer.zero_grad() self.q1_network_optimizer.zero_grad() if self.q2_network is not None: self.q2_network_optimizer.zero_grad() if self.value_network is not None: self.value_network_optimizer.zero_grad() with torch.enable_grad(): # # First, optimize Q networks; minimizing MSE between # Q(s, a) & r + discount * V'(next_s) # current_state_action = rlt.PreprocessedStateAction(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.PreprocessedState(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.PreprocessedState(state=learning_input.next_state)) next_state_actor_action = rlt.PreprocessedStateAction( state=learning_input.next_state, action=rlt.PreprocessedFeatureVector( 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 if self.gamma > 0.0: target_q_value = ( reward + discount * next_state_value * not_done_mask.float()) else: # This is useful in debugging instability issues target_q_value = reward 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.PreprocessedStateAction( state=state, action=rlt.PreprocessedFeatureVector( 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() if self.add_kld_to_loss: if self.apply_kld_on_mean: action_batch_m = torch.mean(actor_output.action_mean, axis=0) action_batch_v = torch.var(actor_output.action_mean, axis=0) else: action_batch_m = torch.mean(actor_output.action, axis=0) action_batch_v = torch.var(actor_output.action, axis=0) kld = (0.5 * ((action_batch_v + (action_batch_m - self.action_emb_mean)**2) / self.action_emb_variance - 1 + self.action_emb_variance.log() - action_batch_v.log()).sum()) actor_loss_mean += self.kld_weight * kld 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 != 0 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) if self.add_kld_to_loss: SummaryWriterContext.add_histogram("kld/mean", action_batch_m) SummaryWriterContext.add_histogram("kld/var", action_batch_v) SummaryWriterContext.add_scalar("kld/kld", kld) 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, )