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,
)
