Graph states

pyproject.toml

@@ -26,6 +26,8 @@ einops = ">=0.6.1"
 numpy = ">=1.21.2"
 python = "^3.10"
 torch = ">=2.6.0"
+tensordict = ">=0.6.1"
+torch_geometric = ">=2.6.1"
 
 # dev dependencies.
 black = { version = "24.3", optional = true }
@@ -99,17 +101,22 @@ include = '\.pyi?$'
 extend-exclude = '''/(\.git|\.hg|\.mypy_cache|\.ipynb|\.tox|\.venv|build)/g'''
 
 [tool.pyright]
+pythonVersion = "3.10"
 include = ["src/gfn", "tutorials/examples", "testing"]  # Removed ** globstars
 exclude = [
     "**/node_modules",
     "**/__pycache__",
     "**/.*", # Exclude dot files and folders
 ]
 
+strict = [
+
+]
+# This is required as the CI pre-commit does not dl the module (i.e. numpy)
+# Therefore, we have to ignore missing imports
 # Removed "strict": [], as it's redundant with typeCheckingMode
 
 typeCheckingMode = "basic"
-pythonVersion = "3.10"
 
 # Removed enableTypeIgnoreComments, not available in pyproject.toml, and bad practice.
 
@@ -124,6 +131,9 @@ reportUntypedFunctionDecorator = "none"
 reportMissingTypeStubs = false
 reportUnboundVariable = "warning"
 reportGeneralTypeIssues = "none"
+reportAttributeAccessIssue = false
+
+[tool.pytest.ini_options]
 reportOptionalMemberAccess = "error"
 reportArgumentType = "error" #This setting doesn't exist, removed.
 

src/gfn/actions.py

@@ -1,10 +1,12 @@
 from __future__ import annotations  # This allows to use the class name in type hints
 
+import enum
 from abc import ABC
 from math import prod
 from typing import ClassVar, List, Sequence
 
 import torch
+from tensordict import TensorDict
 
 
 class Actions(ABC):
@@ -134,7 +136,7 @@ def extend_with_dummy_actions(self, required_first_dim: int) -> None:
                 "extend_with_dummy_actions is only implemented for bi-dimensional actions."
             )
 
-    def compare(self, other: torch.Tensor) -> torch.Tensor:
+    def _compare(self, other: torch.Tensor) -> torch.Tensor:
         """Compares the actions to a tensor of actions.
 
         Args:
@@ -161,12 +163,150 @@ def is_dummy(self) -> torch.Tensor:
         dummy_actions_tensor = self.__class__.dummy_action.repeat(
             *self.batch_shape, *((1,) * len(self.__class__.action_shape))
         )
-        return self.compare(dummy_actions_tensor)
+        return self._compare(dummy_actions_tensor)
 
     @property
     def is_exit(self) -> torch.Tensor:
         """Returns a boolean tensor of shape `batch_shape` indicating whether the actions are exit actions."""
         exit_actions_tensor = self.__class__.exit_action.repeat(
             *self.batch_shape, *((1,) * len(self.__class__.action_shape))
         )
-        return self.compare(exit_actions_tensor)
+        return self._compare(exit_actions_tensor)
+
+
+class GraphActionType(enum.IntEnum):
+    ADD_NODE = 0
+    ADD_EDGE = 1
+    EXIT = 2
+    DUMMY = 3
+
+
+class GraphActions(Actions):
+    """Actions for graph-based environments.
+
+    Each action is one of:
+    - ADD_NODE: Add a node with given features
+    - ADD_EDGE: Add an edge between two nodes with given features
+    - EXIT: Terminate the trajectory
+
+    Attributes:
+        features_dim: Dimension of node/edge features
+        tensor: TensorDict containing:
+            - action_type: Type of action (GraphActionType)
+            - features: Features for nodes/edges
+            - edge_index: Source/target nodes for edges
+    """
+
+    features_dim: ClassVar[int]
+
+    def __init__(self, tensor: TensorDict):
+        """Initializes a GraphAction object.
+
+        Args:
+            action: a GraphActionType indicating the type of action.
+            features: a tensor of shape (batch_shape, feature_shape) representing the features
+                of the nodes or of the edges, depending on the action type. In case of EXIT
+                action, this can be None.
+            edge_index: an tensor of shape (batch_shape, 2) representing the edge to add.
+                This must defined if and only if the action type is GraphActionType.AddEdge.
+        """
+        self.batch_shape = tensor["action_type"].shape
+        features = tensor.get("features", None)
+        if features is None:
+            assert torch.all(
+                torch.logical_or(
+                    tensor["action_type"] == GraphActionType.EXIT,
+                    tensor["action_type"] == GraphActionType.DUMMY,
+                )
+            )
+            features = torch.zeros((*self.batch_shape, self.features_dim))
+        edge_index = tensor.get("edge_index", None)
+        if edge_index is None:
+            assert torch.all(tensor["action_type"] != GraphActionType.ADD_EDGE)
+            edge_index = torch.zeros((*self.batch_shape, 2), dtype=torch.long)
+
+        self.tensor = TensorDict(
+            {
+                "action_type": tensor["action_type"],
+                "features": features,
+                "edge_index": edge_index,
+            },
+            batch_size=self.batch_shape,
+        )
+
+    def __repr__(self):
+        return f"""GraphAction object with {self.batch_shape} actions."""
+
+    def _compare(self, other: GraphActions) -> torch.Tensor:
+        """Compares the actions to another GraphAction object.
+
+        Args:
+            other: GraphAction object to compare.
+
+        Returns: boolean tensor of shape batch_shape indicating whether the actions are equal.
+        """
+        action_compare = torch.all(
+            self.tensor["action_type"] == other.tensor["action_type"]
+        )
+        exit_compare = (
+            torch.all(self.tensor["features"] == other.tensor["features"])
+            | action_compare
+            == GraphActionType.EXIT
+        )
+        edge_compare = (action_compare != GraphActionType.ADD_EDGE) | (
+            torch.all(self.tensor["edge_index"] == other.tensor["edge_index"])
+        )
+        return action_compare & exit_compare & edge_compare
+
+    @property
+    def is_exit(self) -> torch.Tensor:
+        """Returns a boolean tensor of shape `batch_shape` indicating whether the actions are exit actions."""
+        return self.action_type == GraphActionType.EXIT
+
+    @property
+    def is_dummy(self) -> torch.Tensor:
+        """Returns a boolean tensor of shape `batch_shape` indicating whether the actions are dummy actions."""
+        return self.action_type == GraphActionType.DUMMY
+
+    @property
+    def action_type(self) -> torch.Tensor:
+        """Returns the action type tensor."""
+        return self.tensor["action_type"]
+
+    @property
+    def features(self) -> torch.Tensor:
+        """Returns the features tensor."""
+        return self.tensor["features"]
+
+    @property
+    def edge_index(self) -> torch.Tensor:
+        """Returns the edge index tensor."""
+        return self.tensor["edge_index"]
+
+    @classmethod
+    def make_dummy_actions(cls, batch_shape: tuple[int]) -> GraphActions:
+        """Creates a GraphActions object of dummy actions with the given batch shape."""
+        return cls(
+            TensorDict(
+                {
+                    "action_type": torch.full(
+                        batch_shape, fill_value=GraphActionType.DUMMY
+                    ),
+                },
+                batch_size=batch_shape,
+            )
+        )
+
+    @classmethod
+    def make_exit_actions(cls, batch_shape: tuple[int]) -> Actions:
+        """Creates an GraphActions object of exit actions with the given batch shape."""
+        return cls(
+            TensorDict(
+                {
+                    "action_type": torch.full(
+                        batch_shape, fill_value=GraphActionType.EXIT
+                    ),
+                },
+                batch_size=batch_shape,
+            )
+        )

src/gfn/containers/__init__.py

@@ -1,11 +1,11 @@
 from .base import Container
-from .replay_buffer import PrioritizedReplayBuffer, ReplayBuffer
+from .replay_buffer import NormBasedDiversePrioritizedReplayBuffer, ReplayBuffer
 from .state_pairs import StatePairs
 from .trajectories import Trajectories
 from .transitions import Transitions
 
 __all__ = [
-    "PrioritizedReplayBuffer",
+    "NormBasedDiversePrioritizedReplayBuffer",
     "ReplayBuffer",
     "StatePairs",
     "Trajectories",