diff --git a/src/gflownet/algo/trajectory_balance.py b/src/gflownet/algo/trajectory_balance.py index eac57cc6..8cb7c994 100644 --- a/src/gflownet/algo/trajectory_balance.py +++ b/src/gflownet/algo/trajectory_balance.py @@ -276,14 +276,17 @@ def get_idempotent_actions(self, g: Graph, gd: gd.Data, gp: Graph, action: Graph lmask = getattr(gd, action.action.mask_name) nz = lmask.nonzero() # Legal actions are those with a nonzero mask value actions = [iaction if return_aidx else action] + gp_n = len(gp) + gp_e = gp.number_of_edges() for i in nz: aidx = ActionIndex(action_type=iaction[0], row_idx=i[0].item(), col_idx=i[1].item()) if aidx == iaction: continue ga = self.ctx.ActionIndex_to_GraphAction(gd, aidx, fwd=not action.action.is_backward) child = self.env.step(g, ga) - if nx.algorithms.is_isomorphic(child, gp, lambda a, b: a == b, lambda a, b: a == b): - actions.append(aidx if return_aidx else ga) + if len(child) == gp_n and child.number_of_edges() == gp_e: + if nx.algorithms.is_isomorphic(child, gp, lambda a, b: a == b, lambda a, b: a == b): + actions.append(aidx if return_aidx else ga) return actions def construct_batch(self, trajs, cond_info, log_rewards): @@ -774,4 +777,4 @@ def subtb_cum(self, P_F, P_B, F, R, traj_lengths): fr = torch.cat([F[s_idx:e_idx], torch.tensor([R[ep]], device=F.device)]) p = pdiff[s_idx:e_idx] total_loss[ep] = self._loss(subTB(fr, p)).sum() / (n * n + n) * 2 - return total_loss + return total_loss \ No newline at end of file diff --git a/src/gflownet/envs/graph_building_env.py b/src/gflownet/envs/graph_building_env.py index b10b228d..c2dcd9a2 100644 --- a/src/gflownet/envs/graph_building_env.py +++ b/src/gflownet/envs/graph_building_env.py @@ -249,7 +249,7 @@ def parents(self, g: Graph): parents: List[Pair(GraphAction, Graph)] The list of parent-action pairs that lead to `g`. """ - parents: List[Tuple[GraphAction, Graph]] = [] + parents: List[Tuple[GraphAction, Graph, int, int]] = [] # Count node degrees degree: Dict[int, int] = defaultdict(int) for a, b in g.edges: @@ -259,11 +259,14 @@ def parents(self, g: Graph): def add_parent(a, new_g): # Only add parent if the proposed parent `new_g` is not isomorphic # to already identified parents - for ap, gp in parents: + n_nodes = len(new_g) + n_edges = new_g.number_of_edges() + for ap, gp, gp_n, gp_e in parents: # Here we are relying on the dict equality operator for nodes and edges - if is_isomorphic(new_g, gp, lambda a, b: a == b, lambda a, b: a == b): - return - parents.append((a, new_g)) + if n_nodes == gp_n and n_edges == gp_e: + if is_isomorphic(new_g, gp, lambda a, b: a == b, lambda a, b: a == b): + return + parents.append((a, new_g, n_nodes, n_edges)) for a, b in g.edges: if degree[a] > 1 and degree[b] > 1 and len(g.edges[(a, b)]) == 0: @@ -305,7 +308,7 @@ def add_parent(a, new_g): GraphAction(GraphActionType.SetNodeAttr, source=i, attr=k, value=g.nodes[i][k]), graph_without_node_attr(g, i, k), ) - return parents + return [(ap, gp) for ap, gp, _, _ in parents] def count_backward_transitions(self, g: Graph, check_idempotent: bool = False): """Counts the number of parents of g (by default, without checking for isomorphisms)""" @@ -1024,4 +1027,4 @@ def traj_log_n(self, traj): def action_type_to_mask(t: GraphActionType, gbatch: gd.Batch, assert_mask_exists: bool = False): if assert_mask_exists: assert hasattr(gbatch, t.mask_name), f"Mask {t.mask_name} not found in graph data" - return getattr(gbatch, t.mask_name) if hasattr(gbatch, t.mask_name) else torch.ones((1, 1), device=gbatch.x.device) + return getattr(gbatch, t.mask_name) if hasattr(gbatch, t.mask_name) else torch.ones((1, 1), device=gbatch.x.device) \ No newline at end of file