def _get_unmasked_q_values(self, q_network, state: rlt.FeatureData,
slate: rlt.DocList) -> torch.Tensor:
""" Gets the q values from the model and target networks """
batch_size, slate_size, _ = slate.float_features.shape
# TODO: Probably should create a new model type
return q_network(state.repeat_interleave(slate_size, dim=0),
slate.as_feature_data()).view(batch_size, slate_size)
def score(state: rlt.FeatureData) -> torch.Tensor:
tiled_state = state.repeat_interleave(repeats=num_candidates, axis=0)
candidate_docs = state.candidate_docs
assert candidate_docs is not None
actions = candidate_docs.as_feature_data()
q_network.eval()
scores = q_network(tiled_state, actions).view(-1, num_candidates)
q_network.train()
select_prob = F.softmax(candidate_docs.value, dim=1)
assert select_prob.shape == scores.shape
return select_prob * scores