def as_parametric_maxq_training_batch(self):
state_dim = self.states.shape[1]
return rlt.PreprocessedTrainingBatch(
training_input=rlt.PreprocessedParametricDqnInput(
state=rlt.PreprocessedFeatureVector(float_features=self.states),
action=rlt.PreprocessedFeatureVector(float_features=self.actions),
next_state=rlt.PreprocessedFeatureVector(
float_features=self.next_states
),
next_action=rlt.PreprocessedFeatureVector(
float_features=self.next_actions
),
tiled_next_state=rlt.PreprocessedFeatureVector(
float_features=self.possible_next_actions_state_concat[
:, :state_dim
]
),
possible_actions=None,
possible_actions_mask=self.possible_actions_mask,
possible_next_actions=rlt.PreprocessedFeatureVector(
float_features=self.possible_next_actions_state_concat[
:, state_dim:
]
),
possible_next_actions_mask=self.possible_next_actions_mask,
reward=self.rewards,
not_terminal=self.not_terminal,
step=self.step,
time_diff=self.time_diffs,
),
extras=rlt.ExtraData(),
)
def sample_memories(self, batch_size, use_gpu=False, batch_first=False):
"""
:param batch_size: number of samples to return
:param use_gpu: whether to put samples on gpu
:param batch_first: If True, the first dimension of data is batch_size.
If False (default), the first dimension is SEQ_LEN. Therefore,
state's shape is SEQ_LEN x BATCH_SIZE x STATE_DIM, for example. By default,
MDN-RNN consumes data with SEQ_LEN as the first dimension.
"""
sample_indices = np.random.randint(self.memory_size, size=batch_size)
device = torch.device("cuda") if use_gpu else torch.device("cpu")
# state/next state shape: batch_size x seq_len x state_dim
# action shape: batch_size x seq_len x action_dim
# reward/not_terminal shape: batch_size x seq_len
state, action, next_state, reward, not_terminal = map(
lambda x: stack(x).float().to(device),
zip(*self.deque_sample(sample_indices)),
)
if not batch_first:
state, action, next_state, reward, not_terminal = transpose(
state, action, next_state, reward, not_terminal)
training_input = rlt.PreprocessedMemoryNetworkInput(
state=rlt.PreprocessedFeatureVector(float_features=state),
reward=reward,
time_diff=torch.ones_like(reward).float(),
action=action,
next_state=rlt.PreprocessedFeatureVector(
float_features=next_state),
not_terminal=not_terminal,
step=None,
)
return rlt.PreprocessedTrainingBatch(training_input=training_input,
extras=None)
def preprocess_batch(train_batch: Any) -> rlt.PreprocessedTrainingBatch:
obs, action, reward, next_obs, next_action, next_reward, terminal, idxs, possible_actions_mask, log_prob = (
train_batch)
obs = torch.tensor(obs).squeeze(2)
action = torch.tensor(action).float()
reward = torch.tensor(reward).unsqueeze(1)
next_obs = torch.tensor(next_obs).squeeze(2)
next_action = torch.tensor(next_action)
not_terinal = 1.0 - torch.tensor(terminal).unsqueeze(1).float()
idxs = torch.tensor(idxs)
possible_actions_mask = torch.tensor(possible_actions_mask).float()
log_prob = torch.tensor(log_prob)
return rlt.PreprocessedTrainingBatch(
training_input=rlt.PreprocessedPolicyNetworkInput(
state=rlt.PreprocessedFeatureVector(float_features=obs),
action=rlt.PreprocessedFeatureVector(float_features=action),
next_state=rlt.PreprocessedFeatureVector(
float_features=next_obs),
next_action=rlt.PreprocessedFeatureVector(
float_features=next_action),
reward=reward,
not_terminal=not_terinal,
step=None,
time_diff=None,
),
extras=rlt.ExtraData(),
)
def as_slate_q_training_batch(self):
batch_size, state_dim = self.states.shape
action_dim = self.actions.shape[1]
return rlt.PreprocessedTrainingBatch(
training_input=rlt.PreprocessedSlateQInput(
state=rlt.PreprocessedFeatureVector(
float_features=self.states),
next_state=rlt.PreprocessedFeatureVector(
float_features=self.next_states),
tiled_state=rlt.PreprocessedTiledFeatureVector(
float_features=self.
possible_actions_state_concat[:, :state_dim].view(
batch_size, -1, state_dim)),
tiled_next_state=rlt.PreprocessedTiledFeatureVector(
float_features=self.
possible_next_actions_state_concat[:, :state_dim].view(
batch_size, -1, state_dim)),
action=rlt.PreprocessedSlateFeatureVector(
float_features=self.
possible_actions_state_concat[:, state_dim:].view(
batch_size, -1, action_dim),
item_mask=self.possible_actions_mask,
item_probability=self.propensities,
),
next_action=rlt.PreprocessedSlateFeatureVector(
float_features=self.
possible_next_actions_state_concat[:, state_dim:].view(
batch_size, -1, action_dim),
item_mask=self.possible_next_actions_mask,
item_probability=self.next_propensities,
),
reward=self.rewards,
reward_mask=self.rewards_mask,
time_diff=self.time_diffs,
step=self.step,
not_terminal=self.not_terminal,
),
extras=rlt.ExtraData(
mdp_id=self.mdp_ids,
sequence_number=self.sequence_numbers,
action_probability=self.propensities,
max_num_actions=self.max_num_actions,
metrics=self.metrics,
),
)
def as_policy_network_training_batch(self):
return rlt.PreprocessedTrainingBatch(
training_input=rlt.PreprocessedPolicyNetworkInput(
state=rlt.PreprocessedFeatureVector(float_features=self.states),
action=rlt.PreprocessedFeatureVector(float_features=self.actions),
next_state=rlt.PreprocessedFeatureVector(
float_features=self.next_states
),
next_action=rlt.PreprocessedFeatureVector(
float_features=self.next_actions
),
reward=self.rewards,
not_terminal=self.not_terminal,
step=self.step,
time_diff=self.time_diffs,
),
extras=rlt.ExtraData(),
)
def as_cem_training_batch(self, batch_first=False):
"""
Generate one-step samples needed by CEM trainer.
The samples will be used to train an ensemble of world models used by CEM.
If batch_first = True:
state/next state shape: batch_size x 1 x state_dim
action shape: batch_size x 1 x action_dim
reward/terminal shape: batch_size x 1
else (default):
state/next state shape: 1 x batch_size x state_dim
action shape: 1 x batch_size x action_dim
reward/terminal shape: 1 x batch_size
"""
if batch_first:
seq_len_dim = 1
reward, not_terminal = self.rewards, self.not_terminal
else:
seq_len_dim = 0
reward, not_terminal = transpose(self.rewards, self.not_terminal)
training_input = rlt.PreprocessedMemoryNetworkInput(
state=rlt.PreprocessedFeatureVector(
float_features=self.states.unsqueeze(seq_len_dim)),
action=self.actions.unsqueeze(seq_len_dim),
next_state=rlt.PreprocessedFeatureVector(
float_features=self.next_states.unsqueeze(seq_len_dim)),
reward=reward,
not_terminal=not_terminal,
step=self.step,
time_diff=self.time_diffs,
)
return rlt.PreprocessedTrainingBatch(
training_input=training_input,
extras=rlt.ExtraData(
mdp_id=self.mdp_ids,
sequence_number=self.sequence_numbers,
action_probability=self.propensities,
max_num_actions=self.max_num_actions,
metrics=self.metrics,
),
)
def preprocess_batch(train_batch: Any) -> rlt.PreprocessedTrainingBatch:
obs, action, reward, next_obs, next_action, next_reward, terminal, idxs, possible_actions_mask, log_prob = (
train_batch)
batch_size = obs.shape[0]
obs = torch.tensor(obs).squeeze(2)
action = torch.tensor(action).float()
next_obs = torch.tensor(next_obs).squeeze(2)
next_action = torch.tensor(next_action).to(torch.float32)
reward = torch.tensor(reward).unsqueeze(1)
not_terminal = 1 - torch.tensor(terminal).unsqueeze(1).to(torch.uint8)
possible_actions_mask = torch.ones_like(action).to(torch.bool)
tiled_next_state = torch.repeat_interleave(next_obs,
repeats=num_actions,
axis=0)
possible_next_actions = torch.eye(num_actions).repeat(batch_size, 1)
possible_next_actions_mask = not_terminal.repeat(1, num_actions).to(
torch.bool)
return rlt.PreprocessedTrainingBatch(
rlt.PreprocessedParametricDqnInput(
state=rlt.PreprocessedFeatureVector(float_features=obs),
action=rlt.PreprocessedFeatureVector(float_features=action),
next_state=rlt.PreprocessedFeatureVector(
float_features=next_obs),
next_action=rlt.PreprocessedFeatureVector(
float_features=next_action),
possible_actions=None,
possible_actions_mask=possible_actions_mask,
possible_next_actions=rlt.PreprocessedFeatureVector(
float_features=possible_next_actions),
possible_next_actions_mask=possible_next_actions_mask,
tiled_next_state=rlt.PreprocessedFeatureVector(
float_features=tiled_next_state),
reward=reward,
not_terminal=not_terminal,
step=None,
time_diff=None,
),
extras=rlt.ExtraData(),
)
def preprocess_batch(train_batch: Any) -> rlt.PreprocessedTrainingBatch:
obs, action, reward, next_obs, next_action, next_reward, terminal, idxs, possible_actions_mask, log_prob = (
train_batch)
obs = torch.tensor(obs).squeeze(2)
action = torch.tensor(action)
reward = torch.tensor(reward).unsqueeze(1)
next_obs = torch.tensor(next_obs).squeeze(2)
next_action = torch.tensor(next_action)
not_terminal = 1.0 - torch.tensor(terminal).unsqueeze(1).float()
possible_actions_mask = torch.tensor(possible_actions_mask)
next_possible_actions_mask = not_terminal.repeat(1, num_actions)
log_prob = torch.tensor(log_prob)
assert (
action.size(1) == num_actions
), f"action size(1) is {action.size(1)} while num_actions is {num_actions}"
return rlt.PreprocessedTrainingBatch(
training_input=rlt.PreprocessedDiscreteDqnInput(
state=rlt.PreprocessedFeatureVector(float_features=obs),
action=action,
next_state=rlt.PreprocessedFeatureVector(
float_features=next_obs),
next_action=next_action,
possible_actions_mask=possible_actions_mask,
possible_next_actions_mask=next_possible_actions_mask,
reward=reward,
not_terminal=not_terminal,
step=None,
time_diff=None,
),
extras=rlt.ExtraData(
mdp_id=None,
sequence_number=None,
action_probability=log_prob.exp(),
max_num_actions=None,
metrics=None,
),
)
def as_discrete_maxq_training_batch(self):
return rlt.PreprocessedTrainingBatch(
training_input=rlt.PreprocessedDiscreteDqnInput(
state=rlt.PreprocessedFeatureVector(float_features=self.states),
action=self.actions,
next_state=rlt.PreprocessedFeatureVector(
float_features=self.next_states
),
next_action=self.next_actions,
possible_actions_mask=self.possible_actions_mask,
possible_next_actions_mask=self.possible_next_actions_mask,
reward=self.rewards,
not_terminal=self.not_terminal,
step=self.step,
time_diff=self.time_diffs,
),
extras=rlt.ExtraData(
mdp_id=self.mdp_ids,
sequence_number=self.sequence_numbers,
action_probability=self.propensities,
max_num_actions=self.max_num_actions,
metrics=self.metrics,
),
)
def test_seq2slate_eval_data_page(self):
"""
Create 3 slate ranking logs and evaluate using Direct Method, Inverse
Propensity Scores, and Doubly Robust.
The logs are as follows:
state: [1, 0, 0], [0, 1, 0], [0, 0, 1]
indices in logged slates: [3, 2], [3, 2], [3, 2]
model output indices: [2, 3], [3, 2], [2, 3]
logged reward: 4, 5, 7
logged propensities: 0.2, 0.5, 0.4
predicted rewards on logged slates: 2, 4, 6
predicted rewards on model outputted slates: 1, 4, 5
Direct Method uses the predicted rewards on model outputted slates.
Thus the result is expected to be (1 + 4 + 5) / 3
Inverse Propensity Scores would scale the reward by 1.0 / logged propensities
whenever the model output slate matches with the logged slate.
Since only the second log matches with the model output, the IPS result
is expected to be 5 / 0.5 / 3
Doubly Robust is the sum of the direct method result and propensity-scaled
reward difference; the latter is defined as:
1.0 / logged_propensities * (logged reward - predicted reward on logged slate)
* Indicator(model slate == logged slate)
Since only the second logged slate matches with the model outputted slate,
the DR result is expected to be (1 + 4 + 5) / 3 + 1.0 / 0.5 * (5 - 4) / 3
"""
batch_size = 3
state_dim = 3
src_seq_len = 2
tgt_seq_len = 2
candidate_dim = 2
reward_net = FakeSeq2SlateRewardNetwork()
seq2slate_net = FakeSeq2SlateTransformerNet()
baseline_net = nn.Linear(1, 1)
trainer = Seq2SlateTrainer(
seq2slate_net,
baseline_net,
parameters=None,
minibatch_size=3,
use_gpu=False,
)
src_seq = torch.eye(candidate_dim).repeat(batch_size, 1, 1)
tgt_out_idx = torch.LongTensor([[3, 2], [3, 2], [3, 2]])
tgt_out_seq = src_seq[torch.arange(batch_size).
repeat_interleave(tgt_seq_len), # type: ignore
tgt_out_idx.flatten() - 2, ].reshape(
batch_size, tgt_seq_len, candidate_dim)
ptb = rlt.PreprocessedTrainingBatch(
training_input=rlt.PreprocessedRankingInput(
state=rlt.PreprocessedFeatureVector(
float_features=torch.eye(state_dim)),
src_seq=rlt.PreprocessedFeatureVector(float_features=src_seq),
tgt_out_seq=rlt.PreprocessedFeatureVector(
float_features=tgt_out_seq),
src_src_mask=torch.ones(batch_size, src_seq_len, src_seq_len),
tgt_out_idx=tgt_out_idx,
tgt_out_probs=torch.tensor([0.2, 0.5, 0.4]),
slate_reward=torch.tensor([4.0, 5.0, 7.0]),
),
extras=rlt.ExtraData(
sequence_number=torch.tensor([0, 0, 0]),
mdp_id=np.array(["0", "1", "2"]),
),
)
edp = EvaluationDataPage.create_from_training_batch(
ptb, trainer, reward_net)
doubly_robust_estimator = DoublyRobustEstimator()
direct_method, inverse_propensity, doubly_robust = doubly_robust_estimator.estimate(
edp)
logger.info(f"{direct_method}, {inverse_propensity}, {doubly_robust}")
avg_logged_reward = (4 + 5 + 7) / 3
self.assertAlmostEqual(direct_method.raw, (1 + 4 + 5) / 3, delta=1e-6)
self.assertAlmostEqual(direct_method.normalized,
direct_method.raw / avg_logged_reward,
delta=1e-6)
self.assertAlmostEqual(inverse_propensity.raw, 5 / 0.5 / 3, delta=1e-6)
self.assertAlmostEqual(
inverse_propensity.normalized,
inverse_propensity.raw / avg_logged_reward,
delta=1e-6,
)
self.assertAlmostEqual(doubly_robust.raw,
direct_method.raw + 1 / 0.5 * (5 - 4) / 3,
delta=1e-6)
self.assertAlmostEqual(doubly_robust.normalized,
doubly_robust.raw / avg_logged_reward,
delta=1e-6)
