made checklist updates more efficient
@@ -27,17 +27,23 @@ class ChecklistState:
checklist : ``torch.Tensor``
A checklist indicating how many times each action in its agenda has been chosen previously.
It contains the actual counts of the agenda actions.
+ terminal_indices_dict: ``Dict[int, int]``, optional
+ Mapping from batch action indices to indices in any of the four vectors above. If not
+ provided, this mapping will be computed here.
"""
def __init__(self,
terminal_actions: torch.Tensor,
checklist_target: torch.Tensor,
checklist_mask: torch.Tensor,
- checklist: torch.Tensor) -> None:
+ checklist: torch.Tensor,
+ terminal_indices_dict: Dict[int, int] = None) -> None:
self.terminal_actions = terminal_actions
self.checklist_target = checklist_target
self.checklist_mask = checklist_mask
self.checklist = checklist
- # Mapping from batch action indices to indices in any of the four vectors above.
+ if terminal_indices_dict is not None:
+ self.terminal_indices_dict = terminal_indices_dict
+ else:
self.terminal_indices_dict: Dict[int, int] = {}
for checklist_index, batch_action_index in enumerate(terminal_actions.detach().cpu()):
action_index = int(batch_action_index[0])
@@ -54,7 +60,8 @@ class ChecklistState:
new_checklist_state = ChecklistState(terminal_actions=self.terminal_actions,
checklist_target=self.checklist_target,
checklist_mask=self.checklist_mask,
- checklist=new_checklist)
+ checklist=new_checklist,
+ terminal_indices_dict=self.terminal_indices_dict)
return new_checklist_state
def get_balance(self) -> torch.Tensor:
