- probas[1:, :] = pred_probas[:-1, :] * policy[0, :-1, :]
- probas[:-1, :] = pred_probas[1:, :] * policy[1, 1:, :]
- probas[:, 1:] = pred_probas[:, :-1] * policy[2, :, :-1]
- probas[:, :-1] = pred_probas[:, 1:] * policy[3, :, 1:]
- probas[start_i, start_j] = 1.0
+ probas[:, 1:, :] += pred_probas[:, :-1, :] * policies[:, 3, :-1, :]
+ probas[:, :-1, :] += pred_probas[:, 1:, :] * policies[:, 2, 1:, :]
+ probas[:, :, 1:] += pred_probas[:, :, :-1] * policies[:, 1, :, :-1]
+ probas[:, :, :-1] += pred_probas[:, :, 1:] * policies[:, 0, :, 1:]
+ probas[start] = 1.0
+
+ return probas