Repository: AI-Planning/macq (language: python)

Issue:
ARMS (AIJ'07, ICAPS'05)
## Notes
- max-sat based, states+actions

## References
- [x] [Learning action models from plan examples using weighted MAX-SAT](https://www.sciencedirect.com/science/article/pii/S0004370206001408)
- [x] [Learning Action Models from Plan Examples with Incomplete Knowledge](https://www.aaai.org/Papers/ICAPS/2005/ICAPS05-025.pdf)


Discussion / hints:

Fix patch:
diff --git a/.gitignore b/.gitignore
index eae90acb..74d2c9b4 100644
--- a/.gitignore
+++ b/.gitignore
@@ -12,4 +12,5 @@ __init__.pyc
 generated_testing_files/
 new_domain.pddl
 new_prob.pddl
-test_model.json
\ No newline at end of file
+test_model.json
+html/
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index eeda14a5..4025d8a6 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -10,6 +10,8 @@ Install macq for development by cloning the repository and running
 We recommend installing in a virtual environment to avoid package version
 conflicts.
 
+**`tarski` requires [`clingo`](https://potassco.org/clingo/).**
+
 ### Formatting
 
 We use [black](https://black.readthedocs.io/en/stable/) for easy and consistent
diff --git a/README.md b/README.md
index 0bee1b8a..33b212c1 100644
--- a/README.md
+++ b/README.md
@@ -18,7 +18,7 @@ This library is a collection of tools for planning-like action model acquisition
 ## Usage <a name="usage" />
 ```python
 from macq import generate, extract
-from macq.observation import IdentityObservation
+from macq.observation import IdentityObservation, AtomicPartialObservation
 
 # get a domain-specific generator: uses api.planning.domains problem_id/
 # generate 100 traces of length 20 using vanilla sampling
@@ -66,11 +66,12 @@ Model:
         clear location pos-04-04
         at player player-01 location pos-04-06
   ...
+
 ######################################################################
 # Model Extraction - SLAF Technique
 ######################################################################
 traces = generate.pddl.VanillaSampling(problem_id = 123, plan_len = 2, num_traces = 1).traces
-observations = traces.tokenize(PartialObservabilityToken, method=PartialObservabilityToken.random_subset, percent_missing=0.10)
+observations = traces.tokenize(AtomicPartialObservation, percent_missing=0.10)
 model = Extract(observations, modes.SLAF)
 model.details()
 
@@ -92,9 +93,9 @@ Model:
 - [x] [Learning Planning Operators by Observation and Practice](https://aaai.org/Papers/AIPS/1994/AIPS94-057.pdf) (AIPS'94)
 - [ ] [Learning by Experimentation: Incremental Refinement of Incomplete Planning Domains](https://www.sciencedirect.com/science/article/pii/B9781558603356500192) (ICML'94)
 - [ ] [Learning Probabilistic Relational Planning Rules](https://people.csail.mit.edu/lpk/papers/2005/zpk-aaai05.pdf) (ICAPS'04)
-- [ ] [Learning Action Models from Plan Examples with Incomplete Knowledge](https://www.aaai.org/Papers/ICAPS/2005/ICAPS05-025.pdf) (ICAPS'05)
+- [x] [Learning Action Models from Plan Examples with Incomplete Knowledge](https://www.aaai.org/Papers/ICAPS/2005/ICAPS05-025.pdf) (ICAPS'05)
 - [ ] [Learning Planning Rules in Noisy Stochastic Worlds](https://people.csail.mit.edu/lpk/papers/2005/zpk-aaai05.pdf) (AAAI'05)
-- [ ] [Learning action models from plan examples using weighted MAX-SAT](https://www.sciencedirect.com/science/article/pii/S0004370206001408) (AIJ'07)
+- [x] [Learning action models from plan examples using weighted MAX-SAT](https://www.sciencedirect.com/science/article/pii/S0004370206001408) (AIJ'07)
 - [ ] [Learning Symbolic Models of Stochastic Domains](https://www.aaai.org/Papers/JAIR/Vol29/JAIR-2910.pdf) (JAIR'07)
 - [x] [Learning Partially Observable Deterministic Action Models](https://www.aaai.org/Papers/JAIR/Vol33/JAIR-3310.pdf) (JAIR'08)
 - [ ] [Acquisition of Object-Centred Domain Models from Planning Examples](https://ojs.aaai.org/index.php/ICAPS/article/view/13391) (ICAPS'09)
diff --git a/macq/extract/__init__.py b/macq/extract/__init__.py
index 20142941..d5559caf 100644
--- a/macq/extract/__init__.py
+++ b/macq/extract/__init__.py
@@ -1,14 +1,16 @@
-from .model import Model
-from .extract import Extract, modes, IncompatibleObservationToken, SLAF
 from .learned_fluent import LearnedFluent
 from .learned_action import LearnedAction
+from .model import Model, LearnedAction
+from .extract import Extract, modes
+from .exceptions import IncompatibleObservationToken
+from .model import Model
+from .extract import Extract, modes
 
 __all__ = [
+    "LearnedAction",
+    "LearnedFluent",
     "Model",
     "Extract",
     "modes",
     "IncompatibleObservationToken",
-    "LearnedAction",
-    "SLAF",
-    "LearnedFluent",
 ]
diff --git a/macq/extract/arms.py b/macq/extract/arms.py
new file mode 100644
index 00000000..19dd3a01
--- /dev/null
+++ b/macq/extract/arms.py
@@ -0,0 +1,895 @@
+from collections import defaultdict, Counter
+from dataclasses import dataclass
+from warnings import warn
+from typing import Set, List, Dict, Tuple, Hashable
+from nnf import Var, And, Or, false as nnffalse
+from . import LearnedAction, Model, LearnedFluent
+from .exceptions import (
+    IncompatibleObservationToken,
+    InvalidMaxSATModel,
+)
+from ..observation import PartialObservation as Observation, ObservationLists
+from ..trace import Fluent, Action  # Action only used for typing
+from ..utils.pysat import to_wcnf, RC2, WCNF
+
+
+@dataclass
+class Relation:
+    """Fluents with the parameters replaced by their types."""
+
+    name: str
+    types: list
+
+    def var(self):
+        """Generates the variable representation for NNF."""
+        return f"{self.name} {' '.join(list(self.types))}"
+
+    def matches(self, action: LearnedAction):
+        """Determines if a relation is related to a given action."""
+        action_types = set(action.obj_params)
+        self_counts = Counter(self.types)
+        action_counts = Counter(action.obj_params)
+        return all([t in action_types for t in self.types]) and all(
+            [self_counts[t] <= action_counts[t] for t in self.types]
+        )
+
+    def __hash__(self):
+        return hash(self.var())
+
+
+@dataclass
+class ARMSConstraints:
+    """A dataclass to hold all the constraints and weight information."""
+
+    action: List[Or[Var]]
+    info: List[Or[Var]]
+    info3: Dict[Or[Var], int]
+    plan: Dict[Or[Var], int]
+
+
+class ARMS:
+    """ARMS model extraction method.
+
+    Extracts a Model from state observations using the ARMS technique. Fluents
+    are retrieved from the initial state. Actions are learned using the
+    algorithm.
+    """
+
+    class InvalidThreshold(Exception):
+        def __init__(self, threshold):
+            super().__init__(
+                f"Invalid threshold value: {threshold}. Threshold must be a float between 0-1 (inclusive)."
+            )
+
+    def __new__(
+        cls,
+        obs_lists: ObservationLists,
+        debug: bool,
+        upper_bound: int,
+        min_support: int = 2,
+        action_weight: int = 110,
+        info_weight: int = 100,
+        threshold: float = 0.6,
+        info3_default: int = 30,
+        plan_default: int = 30,
+    ):
+        """
+        Arguments:
+            obs_lists (ObservationLists):
+                The observations to extract the model from.
+            upper_bound (int):
+                The upper bound for the maximum size of an action's preconditions and
+                add/delete lists. Determines when an action schemata is fully learned.
+            min_support (int):
+                The minimum support count for an action pair to be considered frequent.
+            action_weight (int):
+                The constant weight W_A(a) to assign to each action constraint.
+                Should be set higher than the weight of information constraints.
+            info_weight (int):
+                The constant weight W_I(r) to assign to each information constraint.
+                Determined empirically, generally the highest in all constraints' weights.
+            threshold (float):
+                (0-1). The probability threshold θ to determine if an I3/plan constraint
+                is weighted by its probability or set to a default value.
+            info3_default (int):
+                The default weight for I3 constraints with probability below the threshold.
+            plan_default (int):
+                The default weight for plan constraints with probability below the threshold.
+        """
+        if obs_lists.type is not Observation:
+            raise IncompatibleObservationToken(obs_lists.type, ARMS)
+
+        if not (threshold >= 0 and threshold <= 1):
+            raise ARMS.InvalidThreshold(threshold)
+
+        fluents = obs_lists.get_fluents()
+        # get fluents from initial state
+        # call algorithm to get actions
+        actions = ARMS._arms(
+            obs_lists,
+            upper_bound,
+            fluents,
+            min_support,
+            action_weight,
+            info_weight,
+            threshold,
+            info3_default,
+            plan_default,
+            debug,
+        )
+
+        learned_fluents = set(map(lambda f: LearnedFluent(f.name, f.objects), fluents))
+        return Model(learned_fluents, actions)
+
+    @staticmethod
+    def _arms(
+        obs_lists: ObservationLists,
+        upper_bound: int,
+        fluents: Set[Fluent],
+        min_support: int,
+        action_weight: int,
+        info_weight: int,
+        threshold: float,
+        info3_default: int,
+        plan_default: int,
+        debug: bool,
+    ) -> Set[LearnedAction]:
+        """The main driver for the ARMS algorithm."""
+        learned_actions = set()  # The set of learned action models Θ
+        # pointers to the earliest unlearned action for each observation list
+        early_actions = [0] * len(obs_lists)
+
+        debug1 = ARMS.debug_menu("Debug step 1?") if debug else False
+        connected_actions, action_map = ARMS.step1(obs_lists, debug1)
+        if debug1:
+            input("Press enter to continue...")
+
+        action_map_rev: Dict[LearnedAction, List[Action]] = defaultdict(list)
+        for obs_action, learned_action in action_map.items():
+            action_map_rev[learned_action].append(obs_action)
+
+        count = 1
+        while action_map_rev:
+            if debug:
+                print("Iteration", count)
+                count += 1
+
+            debug2 = ARMS.debug_menu("Debug step 2?") if debug else False
+            constraints, relation_map = ARMS.step2(
+                obs_lists, connected_actions, action_map, fluents, min_support, debug2
+            )
+            if debug2:
+                input("Press enter to continue...")
+
+            relation_map_rev: Dict[Relation, List[Fluent]] = defaultdict(list)
+            for fluent, relation in relation_map.items():
+                relation_map_rev[relation].append(fluent)
+
+            debug3 = ARMS.debug_menu("Debug step 3?") if debug else False
+            max_sat, decode = ARMS.step3(
+                constraints,
+                action_weight,
+                info_weight,
+                threshold,
+                info3_default,
+                plan_default,
+                debug3,
+            )
+            if debug3:
+                input("Press enter to continue...")
+
+            model = ARMS.step4(max_sat, decode)
+
+            debug5 = ARMS.debug_menu("Debug step 5?") if debug else False
+            # Mutates the LearnedAction (keys) of action_map_rev
+            ARMS.step5(
+                model,
+                list(action_map_rev.keys()),
+                debug5,
+            )
+
+            # Step 5 updates
+
+            # Progress observed states if early actions have been learned
+            setA = set()
+            for action in action_map_rev.keys():
+                for i, obs_list in enumerate(obs_lists):
+                    obs_action: Action = obs_list[early_actions[i]].action
+                    # if current action is the early action for obs_list i,
+                    # update the next state with the effects and update the
+                    # early action pointer
+                    if obs_action in action_map and action == action_map[obs_action]:
+                        print()
+                        # Set add effects true
+                        for add in action.add:
+                            # get candidate fluents from add relation
+                            # get fluent by cross referencing obs_list.action params
+                            candidates = relation_map_rev[add]
+                            for fluent in candidates:
+                                if set(fluent.objects).issubset(obs_action.obj_params):
+                                    obs_list[early_actions[i] + 1].state[fluent] = True
+                                    early_actions[i] += 1
+                        # Set del effects false
+                        for delete in action.delete:
+                            candidates = relation_map_rev[delete]
+                            for fluent in candidates:
+                                if set(fluent.objects).issubset(obs_action.obj_params):
+                                    obs_list[early_actions[i] + 1].state[fluent] = False
+                                    early_actions[i] += 1
+
+                if debug:
+                    print()
+                    print(action.details())
+                    print("precond:", action.precond)
+                    print("add:", action.add)
+                    print("delete:", action.delete)
+
+                if (
+                    max([len(action.precond), len(action.add), len(action.delete)])
+                    >= upper_bound
+                ):
+                    if debug:
+                        print(
+                            f"Action schemata for {action.details()} has been fully learned."
+                        )
+                    setA.add(action)
+
+            # Update Λ by Λ − A
+            for action in setA:
+                action_keys = action_map_rev[action]
+                for obs_action in action_keys:
+                    del action_map[obs_action]
+                del action_map_rev[action]
+                del connected_actions[action]
+                action_keys = [
+                    a1 for a1 in connected_actions if action in connected_actions[a1]
+                ]
+                for a1 in action_keys:
+                    del connected_actions[a1][action]
+
+                # Update Θ by adding A
+                learned_actions.add(action)
+
+            if debug5:
+                input("Press enter to continue...")
+
+        return learned_actions
+
+    @staticmethod
+    def step1(
+        obs_lists: ObservationLists, debug: bool
+    ) -> Tuple[
+        Dict[LearnedAction, Dict[LearnedAction, Set[str]]],
+        Dict[Action, LearnedAction],
+    ]:
+        """(Step 1) Substitute instantiated objects in each action instance with the object type."""
+
+        learned_actions: Set[LearnedAction] = set()
+        action_map: Dict[Action, LearnedAction] = {}
+        for obs_action in obs_lists.get_actions():
+            # We don't support objects with multiple types right now, so no
+            # multiple type clauses need to be generated.
+
+            # Create LearnedActions for each action, replacing instantiated
+            # objects with the object type.
+            types = [obj.obj_type for obj in obs_action.obj_params]
+            learned_action = LearnedAction(obs_action.name, types)
+            learned_actions.add(learned_action)
+            action_map[obs_action] = learned_action
+
+        connected_actions: Dict[LearnedAction, Dict[LearnedAction, Set[str]]] = {}
+        for a1 in learned_actions:
+            connected_actions[a1] = {}
+            for a2 in learned_actions.difference({a1}):  # includes connecting with self
+                intersection = set(a1.obj_params).intersection(set(a2.obj_params))
+                if intersection:
+                    connected_actions[a1][a2] = intersection
+                    if debug:
+                        print(
+                            f"{a1.details()} is connected to {a2.details()} by {intersection}"
+                        )
+
+        return connected_actions, action_map
+
+    @staticmethod
+    def step2(
+        obs_lists: ObservationLists,
+        connected_actions: Dict[LearnedAction, Dict[LearnedAction, Set[str]]],
+        action_map: Dict[Action, LearnedAction],
+        fluents: Set[Fluent],
+        min_support: int,
+        debug: bool,
+    ) -> Tuple[ARMSConstraints, Dict[Fluent, Relation]]:
+        """(Step 2) Generate action constraints, information constraints, and plan constraints.
+
+        For the unexplained actions, build a set of information and action
+        constraints based on individual actions. Apply a frequent-set
+        mining algorithm to find the frequent sets of connected actions and
+        relations. Here connected means the actions and relations must share
+        some common parameters.
+        """
+
+        # Map fluents to relations
+        # relations are fluents but with instantiated objects replaced by the object type
+        relations: Dict[Fluent, Relation] = dict(
+            map(
+                lambda f: (
+                    f,
+                    Relation(
+                        f.name,
+                        [obj.obj_type for obj in f.objects],
+                    ),
+                ),
+                fluents,
+            )
+        )
+
+        debuga = ARMS.debug_menu("Debug action constraints?") if debug else False
+
+        action_constraints = ARMS.step2A(
+            connected_actions, set(relations.values()), debuga
+        )
+
+        debugi = ARMS.debug_menu("Debug info constraints?") if debug else False
+        info_constraints, info_support_counts = ARMS.step2I(
+            obs_lists, relations, action_map, debugi
+        )
+
+        debugp = ARMS.debug_menu("Debug plan constraints?") if debug else False
+        plan_constraints = ARMS.step2P(
+            obs_lists,
+            connected_actions,
+            action_map,
+            set(relations.values()),
+            min_support,
+            debugp,
+        )
+
+        return (
+            ARMSConstraints(
+                action_constraints,
+                info_constraints,
+                info_support_counts,
+                plan_constraints,
+            ),
+            relations,
+        )
+
+    @staticmethod
+    def step2A(
+        connected_actions: Dict[LearnedAction, Dict[LearnedAction, Set]],
+        relations: Set[Relation],
+        debug: bool,
+    ) -> List[Or[Var]]:
+        """Action constraints.
+
+        A1. The intersection of the precondition and add lists of all actions must be empty.
+
+        A2. In addition, if an action’s delete list includes a relation, this relation is
+            in the action’s precondition list. Thus, for every action, we require that
+            the delete list is a subset of the precondition list.
+        """
+
+        if debug:
+            print("\nBuilding action constraints...\n")
+
+        def implication(a: Var, b: Var):
+            return Or([a.negate(), b])
+
+        constraints: List[Or[Var]] = []
+        actions = set(connected_actions.keys())
+        for action in actions:
+            for relation in relations:
+                # A relation is relevant to an action if they share parameter types
+                if relation.matches(action):
+                    if debug:
+                        print(
+                            f'relation ({relation.var()}) is relevant to action "{action.details()}"\n'
+                            "A1:\n"
+                            f"  {relation.var()}∈ add ⇒ {relation.var()}∉ pre\n"
+                            f"  {relation.var()}∈ pre ⇒ {relation.var()}∉ add\n"
+                            "A2:\n"
+                            f"  {relation.var()}∈ del ⇒ {relation.var()}∈ pre\n"
+                        )
+
+                    # A1
+                    # relation in action.add => relation not in action.precond
+                    # relation in action.precond => relation not in action.add
+
+                    # (BREAK) marks unambiguous breakpoints for parsing later
+                    constraints.append(
+                        implication(
+                            Var(
+                                f"{relation.var()} (BREAK) in (BREAK) add (BREAK) {action.details()}"
+                            ),
+                            Var(
+                                f"{relation.var()} (BREAK) in (BREAK) pre (BREAK) {action.details()}"
+                            ).negate(),
+                        )
+                    )
+                    constraints.append(
+                        implication(
+                            Var(
+                                f"{relation.var()} (BREAK) in (BREAK) pre (BREAK) {action.details()}"
+                            ),
+                            Var(
+                                f"{relation.var()} (BREAK) in (BREAK) add (BREAK) {action.details()}"
+                            ).negate(),
+                        )
+                    )
+
+                    # A2
+                    # relation in action.del => relation in action.precond
+                    constraints.append(
+                        implication(
+                            Var(
+                                f"{relation.var()} (BREAK) in (BREAK) del (BREAK) {action.details()}"
+                            ),
+                            Var(
+                                f"{relation.var()} (BREAK) in (BREAK) pre (BREAK) {action.details()}"
+                            ),
+                        )
+                    )
+
+        return constraints
+
+    @staticmethod
+    def step2I(
+        obs_lists: ObservationLists,
+        relations: Dict[Fluent, Relation],
+        actions: Dict[Action, LearnedAction],
+        debug: bool,
+    ) -> Tuple[List[Or[Var]], Dict[Or[Var], int]]:
+        """Information constraints.
+
+        Suppose we observe a relation p to be true between two actions
+        \(a_n\) and \(a_{n+1}\) , and \(p, a_{i_1} , ... ,\) and \(a_{i_k}\) share
+        the same parameter types. We can represent this fact by the following clauses,
+        given that \(a_{i_1} , ... ,\) and \(a_{i_k}\) appear in that order.
+
+        I1. The relation \(p\) must be generated by an action \(a_{i_k} (0 \le i_k \le n)\),
+        that is, \(p\) is selected to be in the add-list of \(a_{i_k}\).
+        \(p∈ (add_{i_1} ∪ add_{i_2} ∪ \dots ∪ add_{i_k}) \), where ∪ means logical “or”.
+
+        I2. The last action \(a_{i_k}\) must not delete the relation p; that is,
+        \(p\) must not be selected to be in the delete list of \(a_{i_k}\): \(p \\not\in del_{i_k}\).
+
+        I3. We define the weight value of a relation-action pair \((p, a)\) as the
+        occurrence probability of this pair in all plan examples. If the probability
+        of a relation-action pair is higher than the probability threshold θ ,
+        then we set a corresponding relation constraint \(p ∈ \\text{PRECOND}_a\), which
+        receives a weight value equal to its prior probability.
+        """
+        if debug:
+            print("\nBuilding information constraints...")
+        constraints: List[Or[Var]] = []
+        support_counts: Dict[Or[Var], int] = defaultdict(int)
+        obs_list: List[Observation]
+        for obs_list_i, obs_list in enumerate(obs_lists):
+            for i, obs in enumerate(obs_list):
+                if obs.state is not None and i > 0:
+                    n = i - 1
+                    if debug:
+                        print(
+                            f"\nStep {i} of observation list {obs_list_i} contains state information."
+                        )
+                    for fluent, val in obs.state.items():
+                        relation = relations[fluent]
+                        # Information constraints only apply to true relations
+                        if val:
+                            if debug:
+                                print(
+                                    f"  Fluent {fluent} is true.\n"
+                                    f"    ({relation.var()})∈ ("
+                                    f"{' ∪ '.join([f'add_{{ {actions[obs_list[ik].action].details()} }}' for ik in range(0,n+1) if obs_list[ik].action in actions] )}"  # type: ignore
+                                    ")"
+                                )
+                            # I1
+                            # relation in the add list of an action <= n (i-1)
+                            i1: List[Var] = []
+                            for obs_i in obs_list[: i - 1]:
+                                if obs_i.action in actions and obs_i.action is not None:
+                                    ai = actions[obs_i.action]
+                                    if relation.matches(ai):
+                                        i1.append(
+                                            Var(
+                                                f"{relation.var()} (BREAK) in (BREAK) add (BREAK) {ai.details()}"
+                                            )
+                                        )
+
+                            # I2
+                            # relation not in del list of action n (i-1)
+                            i2 = None
+                            a_n = obs_list[i - 1].action
+                            if a_n in actions and a_n is not None:
+                                i2 = Var(
+                                    f"{relation.var()} (BREAK) in (BREAK) del (BREAK) {actions[a_n].details()}"
+                                ).negate()
+
+                            if i1:
+                                constraints.append(Or(i1))
+                            if i2:
+                                constraints.append(Or([i2]))
+
+                            # I3
+                            # count occurences
+                            if (
+                                i < len(obs_list) - 1
+                                and obs.action in actions
+                                and obs.action is not None  # for the linter
+                                and relation.matches(actions[obs.action])
+                            ):
+                                # corresponding constraint is related to the current action's precondition list
+                                support_counts[
+                                    Or(
+                                        [
+                                            Var(
+                                                f"{relation.var()} (BREAK) in (BREAK) pre (BREAK) {actions[obs.action].details()}"
+                                            )
+                                        ]
+                                    )
+                                ] += 1
+                            elif (
+                                a_n in actions
+                                and a_n is not None
+                                and relation.matches(actions[a_n])
+                            ):
+                                # corresponding constraint is related to the previous action's add list
+                                support_counts[
+                                    Or(
+                                        [
+                                            Var(
+                                                f"{relation.var()} (BREAK) in (BREAK) add (BREAK) {actions[a_n].details()}"
+                                            )
+                                        ]
+                                    )
+                                ] += 1
+
+        return constraints, support_counts
+
+    @staticmethod
+    def _apriori(
+        action_lists: List[List[LearnedAction]], minsup: int
+    ) -> Dict[Tuple[LearnedAction, LearnedAction], int]:
+        """An implementation of the Apriori algorithm to find frequent ordered pairs of actions."""
+        counts = Counter(
+            [action for action_list in action_lists for action in action_list]
+        )
+        # L1 = {actions that appear >minsup}
+        L1 = set(
+            frozenset([action])
+            for action in filter(lambda k: counts[k] >= minsup, counts.keys())
+        )  # large 1-itemsets
+
+        # Only going up to L2, so no loop or generalized algorithm needed
+        # apriori-gen step
+        C2 = set([i.union(j) for i in L1 for j in L1 if len(i.union(j)) == 2])
+        # Since L1 contains 1-itemsets where each item is frequent, C2 can
+        # only contain valid sets and pruning is not required
+
+        # Get all possible ordered action pairs
+        C2_ordered = set()
+        for pair in C2:
+            pair = list(pair)
+            C2_ordered.add((pair[0], pair[1]))
+            C2_ordered.add((pair[1], pair[0]))
+
+        # Count pair occurences and generate L2
+        frequent_pairs = {}
+        for ai, aj in C2_ordered:
+            count = 0
+            for action_list in action_lists:
+                a1_indecies = [i for i, e in enumerate(action_list) if e == ai]
+                if a1_indecies:
+                    for i in a1_indecies:
+                        if aj in action_list[i + 1 :]:
+                            count += 1
+            if count >= minsup:
+                frequent_pairs[(ai, aj)] = count
+
+        return frequent_pairs
+
+    @staticmethod
+    def step2P(
+        obs_lists: ObservationLists,
+        connected_actions: Dict[LearnedAction, Dict[LearnedAction, Set]],
+        action_map: Dict[Action, LearnedAction],
+        relations: Set[Relation],
+        min_support: int,
+        debug: bool,
+    ) -> Dict[Or[Var], int]:
+        """Plan constraints.
+
+        P1. Every precondition \(p\) of every action \(b\) must be in the add
+        list of a preceding action \(a\) and is not deleted by any actions between
+        \(a\) and \(b\).
+
+        P2. In addition, at least one relation \(r\) in the add list of an action
+        must be useful in achieving a precondition of a later action. That is, for
+        every action \(a\), an add list relation \(r\) must be in the precondition of a
+        later action \(b\), and there is no other action between \(a\) and \(b\)
+        that either adds or deletes \(r\).
+
+        "While constraints P1 and P2 provide the general guiding principle for
+        ensuring a plan’s correctness, in practice there are too many
+        instantiations of these constraints." (more information on the
+        rationelle in the paper) Thus, they are replaced with the following
+        constraints:
+
+        Let there be an action pair \(\langle a_i, a_j \\rangle, 0 \le i < j \le n-1\).
+
+        P3. One of the relevant relations \(p\) must be chosen to be in the
+        preconditions of both \(a_i\) and \(a_j\), but not in the delete list
+        of \(a_i\).
+
+        P4. The first action \(a_i\) adds a relevant relation that is in the
+        precondition list of the second action \(a_j\) in the pair.
+
+        P5. A relevant relation \(p\) that is deleted by the first action
+        \(a_i\) is added by \(a_j\). The second clause is designed for the event
+        when an action re-establishes a fact that is deleted by a previous action.
+
+        The above constraints can be combined into one constraint:
+        $$\exists p ((p \in (pre_i \cap pre_j) \land p \\not \in (del_i)) \lor (p \in (add_i \cap pre_j)) \lor (p \in (del_i \cap add_j)))$$
+        """
+        frequent_pairs = ARMS._apriori(
+            [
+                [
+                    action_map[obs.action]
+                    for obs in obs_list
+                    if obs.action is not None and obs.action in action_map
+                ]
+                for obs_list in obs_lists
+            ],
+            min_support,
+        )
+        if debug:
+            print("Frequent pairs:")
+            print(frequent_pairs)
+
+        # constraints: Dict[And[Or[Var]], int] = {}
+        constraints: Dict[Or[Var], int] = {}
+        for ai, aj in frequent_pairs.keys():
+            connectors = set()
+            # get list of relevant relations from connected_actions
+            if ai in connected_actions and aj in connected_actions[ai]:
+                connectors.update(connected_actions[ai][aj])
+            if aj in connected_actions and ai in connected_actions[aj]:
+                connectors.update(connected_actions[aj][ai])
+
+            # if the actions are not related they are not a valid pair for a plan constraint.
+            if not connectors:
+                continue
+
+            # for each relation, save constraint
+            relevant_relations = {p for p in relations if connectors.issubset(p.types)}
+            # relation_constraints: List[Or[And[Var]]] = []
+            relation_constraints: List[Var] = []
+            for relation in relevant_relations:
+
+                relation_constraints.append(
+                    Var(
+                        f"{relation.var()} (BREAK) relevant (BREAK) {ai.details()} (BREAK) {aj.details()}"
+                    )
+                )
+                if debug:
+                    print(
+                        f"{relation.var()} might explain action pair ({ai.details()}, {aj.details()})"
+                    )
+            constraints[Or(relation_constraints)] = frequent_pairs[(ai, aj)]
+
+        return constraints
+
+    @staticmethod
+    def step3(
+        constraints: ARMSConstraints,
+        action_weight: int,
+        info_weight: int,
+        threshold: float,
+        info3_default: int,
+        plan_default: int,
+        debug: bool,
+    ) -> Tuple[WCNF, Dict[int, Hashable]]:
+        """(Step 3) Construct the weighted MAX-SAT problem based on the constraints and weight information found in Step 2."""
+
+        action_weights = [action_weight] * len(constraints.action)
+        info_weights = [info_weight] * len(constraints.info)
+        info3_weights = ARMS._calculate_support_rates(
+            list(constraints.info3.values()), threshold, info3_default
+        )
+        plan_weights = ARMS._calculate_support_rates(
+            list(constraints.plan.values()), threshold, plan_default
+        )
+        all_weights = action_weights + info_weights + info3_weights + plan_weights
+
+        info3_constraints = list(constraints.info3.keys())
+        plan_constraints = list(constraints.plan.keys())
+        all_constraints = (
+            constraints.action + constraints.info + info3_constraints + plan_constraints
+        )
+
+        constraints_w_weights: Dict[Or[Var], int] = {}
+        for constraint, weight in zip(all_constraints, all_weights):
+            if constraint == nnffalse:
+                continue
+            if constraint not in constraints_w_weights:
+                constraints_w_weights[constraint] = weight
+            elif weight != constraints_w_weights[constraint]:
+                if debug:
+                    warn(
+                        f"The constraint {constraint} has conflicting weights ({weight} and {constraints_w_weights[constraint]}). Choosing the smaller weight."
+                    )
+                constraints_w_weights[constraint] = min(
+                    weight, constraints_w_weights[constraint]
+                )
+
+        problem: And[Or[Var]] = And(list(constraints_w_weights.keys()))
+        weights = list(constraints_w_weights.values())
+
+        wcnf, decode = to_wcnf(problem, weights)
+        return wcnf, decode
+
+    @staticmethod
+    def _calculate_support_rates(
+        support_counts: List[int], threshold: float, default: int
+    ) -> List[int]:
+        # NOTE:
+        # In the paper, Z_Σ_P (denominator of the support rate formula) is
+        # defined as the "total pairs" in the set of plans. However, in the
+        # examples it appears that they use the max support count as the
+        # denominator. My best interpretation is then to use the max support
+        # count as the denominator to calculate the support rate.
+
+        z_sigma_p = max(support_counts)
+
+        def get_support_rate(count):
+            probability = count / z_sigma_p
+            return probability * 100 if probability > threshold else default
+
+        return list(map(get_support_rate, support_counts))
+
+    @staticmethod
+    def step4(max_sat: WCNF, decode: Dict[int, Hashable]) -> Dict[Hashable, bool]:
+        """(Step 4) Solve the MAX-SAT problem built in Step 3."""
+        solver = RC2(max_sat)
+
+        encoded_model = solver.compute()
+
+        if not isinstance(encoded_model, list):
+            # should never be reached
+            raise InvalidMaxSATModel(encoded_model)
+
+        # decode the model (back to nnf vars)
+        model: Dict[Hashable, bool] = {
+            decode[abs(clause)]: clause > 0 for clause in encoded_model
+        }
+
+        return model
+
+    @staticmethod
+    def step5(
+        model: Dict[Hashable, bool],
+        actions: List[LearnedAction],
+        debug: bool,
+    ):
+        """(Step 5) Extract the learned action effects from the solved model."""
+        action_map = {a.details(): a for a in actions}
+        negative_constraints = defaultdict(set)
+        plan_constraints: List[Tuple[str, LearnedAction, LearnedAction]] = []
+
+        # NOTE: only taking the top n (optimal number varies, determine
+        # empirically) constraints usually results in more accurate action
+        # models, however this is not a part of the paper and therefore not
+        # implemented.
+        for constraint, val in model.items():
+            constraint = str(constraint).split(" (BREAK) ")
+            relation = constraint[0]
+            ctype = constraint[1]  # constraint type
+            if ctype == "in":
+                effect = constraint[2]
+                action = action_map[constraint[3]]
+                if debug:
+                    print(
+                        f"Learned constraint: {relation} in {effect}_{action.details()}"
+                    )
+                if val:
+                    action_update = {
+                        "pre": action.update_precond,
+                        "add": action.update_add,
+                        "del": action.update_delete,
+                    }[effect]
+                    action_update({relation})
+                else:
+                    action_effect = {
+                        "pre": action.precond,
+                        "add": action.add,
+                        "del": action.delete,
+                    }[effect]
+                    if relation in action_effect:
+                        if debug:
+                            warn(
+                                f"Removing {relation} from {effect} of {action.details()}"
+                            )
+                        action_effect.remove(relation)
+                    negative_constraints[(relation, action)].add(effect)
+
+            else:  # store plan constraint
+                ai = action_map[constraint[2]]
+                aj = action_map[constraint[3]]
+                plan_constraints.append((relation, ai, aj))
+                if debug:
+                    print(f"{relation} possibly explains action pair ({ai}, {aj})")
+
+        for p, ai, aj in plan_constraints:
+            if (
+                not (
+                    p in ai.precond.intersection(aj.precond) and p not in ai.delete
+                )  # P3
+                or not (p in ai.add.intersection(aj.precond))  # P4
+                or not (p in ai.delete.intersection(aj.add))  # P5
+            ):
+                # check if either P3 or P4 are partially fulfilled and can be satisfied
+                if p in ai.precond.union(aj.precond):
+                    if p in aj.precond:
+                        # if P3 isn't contradicted, add p to ai.precond
+                        if p not in ai.delete and not (
+                            (p, ai) in negative_constraints
+                            and "pre" in negative_constraints[(p, ai)]
+                        ):
+                            ai.update_precond({p})
+
+                        # if P4 isn't contradicted, add p to ai.add
+                        if not (
+                            (p, ai) in negative_constraints
+                            and "add" in negative_constraints[(p, ai)]
+                        ):
+                            ai.update_add({p})
+
+                    # p in ai.precond and P3 not contradicted, add p to aj.precond
+                    elif p not in ai.delete and not (
+                        (p, aj) in negative_constraints
+                        and "pre" in negative_constraints[(p, aj)]
+                    ):
+                        aj.update_precond({p})
+
+                # check if either P3 or P4 can be satisfied
+                elif not (
+                    (p, aj) in negative_constraints
+                    and "pre" in negative_constraints[(p, aj)]
+                ):
+                    # if P3 isn't contradicted, add p to both ai and aj preconds
+                    if p not in ai.delete and not (
+                        (p, ai) in negative_constraints
+                        and "pre" in negative_constraints[(p, ai)]
+                    ):
+                        ai.update_precond({p})
+                        aj.update_precond({p})
+
+                    # if P4 isn't contradicted, add p to ai.add and aj.precond
+                    if not (
+                        (p, ai) in negative_constraints
+                        and "add" in negative_constraints[(p, ai)]
+                    ):
+                        ai.update_add({p})
+                        aj.update_precond({p})
+
+                # check if P5 can be satisfied
+                # if P5 isn't contradicted, add p wherever it is missing
+                if not (
+                    (p, ai) in negative_constraints
+                    and "del" in negative_constraints[(p, ai)]
+                ) and not (
+                    (p, aj) in negative_constraints
+                    and "add" in negative_constraints[(p, aj)]
+                ):
+                    if p not in ai.delete:
+                        ai.update_delete({p})
+                    if p not in aj.add:
+                        aj.update_add({p})
+
+    @staticmethod
+    def debug_menu(prompt: str):
+        choice = input(prompt + " (y/n): ").lower()
+        while choice not in ["y", "n"]:
+            choice = input(prompt + " (y/n): ").lower()
+        return choice == "y"
diff --git a/macq/extract/exceptions.py b/macq/extract/exceptions.py
new file mode 100644
index 00000000..7407cb23
--- /dev/null
+++ b/macq/extract/exceptions.py
@@ -0,0 +1,26 @@
+class IncompatibleObservationToken(Exception):
+    def __init__(self, token, technique, message=None):
+        if message is None:
+            message = f"Observations of type {token.__name__} are not compatible with the {technique.__name__} extraction technique."
+        super().__init__(message)
+
+
+class InconsistentConstraintWeights(Exception):
+    def __init__(self, constraint, weight1, weight2, message=None):
+        if message is None:
+            message = f"Tried to assign the constraint {constraint} conflicting weights ({weight1} and {weight2})."
+        super().__init__(message)
+
+
+class InvalidMaxSATModel(Exception):
+    def __init__(self, model, message=None):
+        if message is None:
+            message = f"The MAX-SAT solver generated an invalid model. Model should be a list of integers. model = {model}."
+        super().__init__(message)
+
+
+class ConstraintContradiction(Exception):
+    def __init__(self, relation, effect, action, message=None):
+        if message is None:
+            message = f"Action model has contradictory constraints for {relation}'s presence in the {effect} list of {action.details()}."
+        super().__init__(message)
diff --git a/macq/extract/extract.py b/macq/extract/extract.py
index 70e1e1d7..39def5bd 100644
--- a/macq/extract/extract.py
+++ b/macq/extract/extract.py
@@ -5,6 +5,7 @@
 from .model import Model
 from .observer import Observer
 from .slaf import SLAF
+from .arms import ARMS
 
 
 @dataclass
@@ -14,13 +15,6 @@ class SAS:
     post_state: State
 
 
-class IncompatibleObservationToken(Exception):
-    def __init__(self, token, technique, message=None):
-        if message is None:
-            message = f"Observations of type {token.__name__} are not compatible with the {technique.__name__} extraction technique."
-        super().__init__(message)
-
-
 class modes(Enum):
     """Model extraction techniques.
 
@@ -29,6 +23,7 @@ class modes(Enum):
 
     OBSERVER = auto()
     SLAF = auto()
+    ARMS = auto()
 
 
 class Extract:
@@ -38,7 +33,9 @@ class Extract:
     from state observations.
     """
 
-    def __new__(cls, obs_lists: ObservationLists, mode: modes, **kwargs) -> Model:
+    def __new__(
+        cls, obs_lists: ObservationLists, mode: modes, debug: bool = False, **kwargs
+    ) -> Model:
         """Extracts a Model object.
 
         Extracts a model from the observations using the specified extraction
@@ -59,11 +56,10 @@ def __new__(cls, obs_lists: ObservationLists, mode: modes, **kwargs) -> Model:
         techniques = {
             modes.OBSERVER: Observer,
             modes.SLAF: SLAF,
+            modes.ARMS: ARMS,
         }
         if mode == modes.SLAF:
-            # only allow one trace
-            assert (
-                len(obs_lists) == 1
-            ), "The SLAF extraction technique only takes one trace."
+            if len(obs_lists) != 1:
+                raise Exception("The SLAF extraction technique only takes one trace.")
 
-        return techniques[mode](obs_lists, **kwargs)
+        return techniques[mode](obs_lists, debug, **kwargs)
diff --git a/macq/extract/learned_action.py b/macq/extract/learned_action.py
index 89b7ae47..6d040b89 100644
--- a/macq/extract/learned_action.py
+++ b/macq/extract/learned_action.py
@@ -1,9 +1,9 @@
 from __future__ import annotations
-from typing import Set, List
+from typing import List, Set
 
 
 class LearnedAction:
-    def __init__(self, name: str, obj_params: List, **kwargs):
+    def __init__(self, name: str, obj_params: List[str], **kwargs):
         self.name = name
         self.obj_params = obj_params
         if "cost" in kwargs:
@@ -27,7 +27,7 @@ def __hash__(self):
     def details(self):
         # obj_params can be either a list of strings or a list of PlanningObject depending on the token type and extraction method used to learn the action
         try:
-            string = f"({self.name} {' '.join([o for o in self.obj_params])})"
+            string = f"({self.name} {' '.join(self.obj_params)})"
         except TypeError:
             string = f"({self.name} {' '.join([o.details() for o in self.obj_params])})"
 
@@ -60,6 +60,11 @@ def update_delete(self, fluents: Set[str]):
         """
         self.delete.update(fluents)
 
+    def clear(self):
+        self.precond = set()
+        self.add = set()
+        self.delete = set()
+
     def compare(self, orig_action: LearnedAction):
         """Compares the learned action to an original, ground truth action."""
         precond_diff = orig_action.precond.difference(self.precond)
diff --git a/macq/extract/observer.py b/macq/extract/observer.py
index c0aff35f..3647282c 100644
--- a/macq/extract/observer.py
+++ b/macq/extract/observer.py
@@ -2,7 +2,8 @@
 from collections import defaultdict
 
 from attr import dataclass
-import macq.extract as extract
+from . import LearnedAction, Model
+from .exceptions import IncompatibleObservationToken
 from .model import Model
 from .learned_fluent import LearnedFluent
 from ..observation import IdentityObservation, ObservationLists
@@ -29,7 +30,7 @@ class Observer:
     fluents that went from True to False.
     """
 
-    def __new__(cls, obs_lists: ObservationLists):
+    def __new__(cls, obs_lists: ObservationLists, debug: bool):
         """Creates a new Model object.
 
         Args:
@@ -40,7 +41,7 @@ def __new__(cls, obs_lists: ObservationLists):
                 Raised if the observations are not identity observation.
         """
         if obs_lists.type is not IdentityObservation:
-            raise extract.IncompatibleObservationToken(obs_lists.type, Observer)
+            raise IncompatibleObservationToken(obs_lists.type, Observer)
         fluents = Observer._get_fluents(obs_lists)
         actions = Observer._get_actions(obs_lists)
         return Model(fluents, actions)
@@ -76,7 +77,7 @@ def _get_actions(obs_lists: ObservationLists):
         action_transitions = obs_lists.get_all_transitions()
         for action, transitions in action_transitions.items():
             # Create a LearnedAction for the current action
-            model_action = extract.LearnedAction(
+            model_action = LearnedAction(
                 action.name, action.obj_params, cost=action.cost
             )
 
diff --git a/macq/generate/pddl/generator.py b/macq/generate/pddl/generator.py
index 95250395..aea73c12 100644
--- a/macq/generate/pddl/generator.py
+++ b/macq/generate/pddl/generator.py
@@ -1,3 +1,4 @@
+from time import sleep
 from typing import Set, List, Union
 from tarski.io import PDDLReader
 from tarski.search import GroundForwardSearchModel
@@ -21,15 +22,21 @@
 from ...trace import Action, State, PlanningObject, Fluent, Trace, Step
 
 
+class PlanningDomainsAPIError(Exception):
+    """Raised when a valid response cannot be obtained from the planning.domains solver."""
+
+    def __init__(self, message):
+        super().__init__(message)
+
+
 class InvalidGoalFluent(Exception):
     """
     Raised when the user attempts to supply a new goal with invalid fluent(s).
     """
 
-    def __init__(
-        self,
-        message="The fluents provided contain one or more fluents not available in this problem.",
-    ):
+    def __init__(self, fluent, message=None):
+        if message is None:
+            message = f"{fluent} is not available in this problem."
         super().__init__(message)
 
 
@@ -57,10 +64,16 @@ class Generator:
         op_dict (dict):
             The problem's ground operators, formatted to a dictionary for easy access during plan generation.
         observe_pres_effs (bool):
-            Option to observe action preconditions and effects upon generation. 
+            Option to observe action preconditions and effects upon generation.
     """
 
-    def __init__(self, dom: str = None, prob: str = None, problem_id: int = None, observe_pres_effs: bool = False):
+    def __init__(
+        self,
+        dom: str = None,
+        prob: str = None,
+        problem_id: int = None,
+        observe_pres_effs: bool = False,
+    ):
         """Creates a basic PDDL state trace generator. Takes either the raw filenames
         of the domain and problem, or a problem ID.
 
@@ -72,7 +85,7 @@ def __init__(self, dom: str = None, prob: str = None, problem_id: int = None, ob
             problem_id (int):
                 The ID of the problem to access.
             observe_pres_effs (bool):
-                Option to observe action preconditions and effects upon generation. 
+                Option to observe action preconditions and effects upon generation.
         """
         # get attributes
         self.pddl_dom = dom
@@ -248,9 +261,7 @@ def tarski_act_to_macq(self, tarski_act: PlainOperator):
             raw_precond = tarski_act.precondition.subformulas
             for raw_p in raw_precond:
                 if isinstance(raw_p, CompoundFormula):
-                    precond.add(
-                        self.__tarski_atom_to_macq_fluent(raw_p.subformulas[0])
-                    )
+                    precond.add(self.__tarski_atom_to_macq_fluent(raw_p.subformulas[0]))
                 else:
                     precond.add(self.__tarski_atom_to_macq_fluent(raw_p))
         else:
@@ -263,7 +274,17 @@ def tarski_act_to_macq(self, tarski_act: PlainOperator):
         for fluent in precond:
             objs.update(set(fluent.objects))
 
-        return Action(name=name, obj_params=list(objs), precond=precond, add=add, delete=delete) if self.observe_pres_effs else Action(name=name, obj_params=list(objs))
+        return (
+            Action(
+                name=name,
+                obj_params=list(objs),
+                precond=precond,
+                add=add,
+                delete=delete,
+            )
+            if self.observe_pres_effs
+            else Action(name=name, obj_params=list(objs))
+        )
 
     def change_init(
         self,
@@ -285,7 +306,10 @@ def change_init(
         init = create(self.lang)
         for f in init_fluents:
             # convert fluents to tarski Atoms
-            atom = Atom(self.lang.get_predicate(f.name), [self.lang.get(o.name) for o in f.objects])
+            atom = Atom(
+                self.lang.get_predicate(f.name),
+                [self.lang.get(o.name) for o in f.objects],
+            )
             init.add(atom.predicate, *atom.subterms)
         self.problem.init = init
 
@@ -320,7 +344,7 @@ def change_goal(
         available_f = self.grounded_fluents
         for f in goal_fluents:
             if f not in available_f:
-                raise InvalidGoalFluent()
+                raise InvalidGoalFluent(f)
 
         # convert the given set of fluents into a formula
         if not goal_fluents:
@@ -344,7 +368,7 @@ def change_goal(
         self.pddl_dom = new_domain
         self.pddl_prob = new_prob
 
-    def generate_plan(self, from_ipc_file:bool=False, filename:str=None):
+    def generate_plan(self, from_ipc_file: bool = False, filename: str = None):
         """Generates a plan. If reading from an IPC file, the `Plan` is read directly. Otherwise, if the initial state or
         goal was changed, these changes are taken into account through the updated PDDL files. If no changes were made, the
         default nitial state/goal in the initial problem file is used.
@@ -369,14 +393,30 @@ def generate_plan(self, from_ipc_file:bool=False, filename:str=None):
                     "domain": open(self.pddl_dom, "r").read(),
                     "problem": open(self.pddl_prob, "r").read(),
                 }
-                resp = requests.post(
-                    "http://solver.planning.domains/solve", verify=False, json=data
-                ).json()
-                plan = [act["name"] for act in resp["result"]["plan"]]
+
+                def get_api_response(delays: List[int]):
+                    if delays:
+                        sleep(delays[0])
+                        try:
+                            resp = requests.post(
+                                "http://solver.planning.domains/solve",
+                                verify=False,
+                                json=data,
+                            ).json()
+                            return [act["name"] for act in resp["result"]["plan"]]
+                        except TypeError:
+                            return get_api_response(delays[1:])
+
+                plan = get_api_response([0, 1, 3, 5, 10])
+                if plan is None:
+                    raise PlanningDomainsAPIError(
+                        "Could not get a valid response from the planning.domains solver after 5 attempts.",
+                    )
+
         else:
             f = open(filename, "r")
-            plan = list(filter(lambda x: ';' not in x, f.read().splitlines()))
-            
+            plan = list(filter(lambda x: ";" not in x, f.read().splitlines()))
+
         # convert to a list of tarski PlainOperators (actions)
         return Plan([self.op_dict[p] for p in plan if p in self.op_dict])
 
@@ -406,4 +446,4 @@ def generate_single_trace_from_plan(self, plan: Plan):
                 state = progress(state, act)
             else:
                 trace.append(Step(macq_state, None, i + 1))
-        return trace
\ No newline at end of file
+        return trace
diff --git a/macq/generate/pddl/random_goal_sampling.py b/macq/generate/pddl/random_goal_sampling.py
index 17ea53e2..c1c9f210 100644
--- a/macq/generate/pddl/random_goal_sampling.py
+++ b/macq/generate/pddl/random_goal_sampling.py
@@ -8,14 +8,13 @@
 from ...utils.timer import basic_timer
 
 
-
 MAX_GOAL_SEARCH_TIME = 30.0
 
 
 class RandomGoalSampling(VanillaSampling):
     """Random Goal State Trace Sampler - inherits the VanillaSampling class and its attributes.
 
-    A state trace generator that generates traces by randomly generating some candidate states/goals k steps deep, 
+    A state trace generator that generates traces by randomly generating some candidate states/goals k steps deep,
     then running a planner on a random subset of the fluents to get plans. The longest plans (those closest to k, thus representing
     goal states that are somewhat complex and take longer to reach) are taken and used to generate traces.
 
@@ -30,8 +29,9 @@ class RandomGoalSampling(VanillaSampling):
             The percentage of fluents to extract to use as a goal state from the generated states.
         goals_inits_plans (List[Dict]):
             A list of dictionaries, where each dictionary stores the generated goal state as the key and the initial state and plan used to
-            reach the goal as values. 
+            reach the goal as values.
     """
+
     def __init__(
         self,
         steps_deep: int,
@@ -65,20 +65,26 @@ def __init__(
             problem_id (int):
                 The ID of the problem to access.
             observe_pres_effs (bool):
-                Option to observe action preconditions and effects upon generation. 
+                Option to observe action preconditions and effects upon generation.
         """
         if subset_size_perc < 0 or subset_size_perc > 1:
             raise PercentError()
-        self.steps_deep = steps_deep        
+        self.steps_deep = steps_deep
         self.enforced_hill_climbing_sampling = enforced_hill_climbing_sampling
         self.subset_size_perc = subset_size_perc
         self.goals_inits_plans = []
-        super().__init__(dom=dom, prob=prob, problem_id=problem_id, observe_pres_effs=observe_pres_effs, num_traces=num_traces)
+        super().__init__(
+            dom=dom,
+            prob=prob,
+            problem_id=problem_id,
+            observe_pres_effs=observe_pres_effs,
+            num_traces=num_traces,
+        )
 
     def goal_sampling(self):
         """Samples goals by randomly generating candidate goal states k (`steps_deep`) steps deep, then running planners on those
-        goal states to ensure the goals are complex enough (i.e. cannot be reached in too few steps). Candidate 
-        goal states are generated for a set amount of time indicated by MAX_GOAL_SEARCH_TIME, and the goals with the 
+        goal states to ensure the goals are complex enough (i.e. cannot be reached in too few steps). Candidate
+        goal states are generated for a set amount of time indicated by MAX_GOAL_SEARCH_TIME, and the goals with the
         longest plans (the most complex goals) are selected.
 
         Returns: An OrderedDict holding the longest goal states along with the initial state and plans used to reach them.
@@ -86,8 +92,10 @@ def goal_sampling(self):
         goal_states = {}
         self.generate_goals(goal_states=goal_states)
         # sort the results by plan length and get the k largest ones
-        filtered_goals = OrderedDict(sorted(goal_states.items(), key=lambda x : len(x[1]["plan"].actions)))
-        to_del = list(filtered_goals.keys())[:len(filtered_goals) - self.num_traces]
+        filtered_goals = OrderedDict(
+            sorted(goal_states.items(), key=lambda x: len(x[1]["plan"].actions))
+        )
+        to_del = list(filtered_goals.keys())[: len(filtered_goals) - self.num_traces]
         for d in to_del:
             del filtered_goals[d]
         return filtered_goals
@@ -122,12 +130,8 @@ def generate_goals(self, goal_states: Dict):
             self.change_goal(goal_fluents=goal_f)
 
             # ensure that the goal doesn't hold in the initial state; restart if it does
-            init_state = {
-                str(a) for a in self.problem.init.as_atoms()
-            }
-            goal = {
-                str(a) for a in self.problem.goal.subformulas
-            }
+            init_state = {str(a) for a in self.problem.init.as_atoms()}
+            goal = {str(a) for a in self.problem.goal.subformulas}
 
             if goal.issubset(init_state):
                 continue
@@ -144,7 +148,10 @@ def generate_goals(self, goal_states: Dict):
             for f in goal_f:
                 state_dict[f] = True
             # map each goal to the initial state and plan used to achieve it
-            goal_states[State(state_dict)] = {"plan": test_plan, "initial state": self.problem.init}
+            goal_states[State(state_dict)] = {
+                "plan": test_plan,
+                "initial state": self.problem.init,
+            }
 
             # optionally change the initial state of the sampler for the next iteration to the goal state just generated (ensures more diversity in goals/plans)
             # use the full state the goal was extracted from as the initial state to prevent planning errors from incomplete initial states
@@ -155,8 +162,8 @@ def generate_goals(self, goal_states: Dict):
             if len(test_plan.actions) >= self.steps_deep:
                 k_length_plans += 1
             if k_length_plans >= self.num_traces:
-                 break
-            
+                break
+
     def generate_traces(self):
         """Generates traces based on the sampled goals. Traces are generated using the initial state and plan used to achieve the goal.
 
@@ -172,7 +179,5 @@ def generate_traces(self):
             if self.enforced_hill_climbing_sampling:
                 self.problem.init = goal["initial state"]
             # generate a plan based on the new goal/initial state, then generate a trace based on that plan
-            traces.append(
-                self.generate_single_trace_from_plan(goal["plan"])
-            )
-        return traces
\ No newline at end of file
+            traces.append(self.generate_single_trace_from_plan(goal["plan"]))
+        return traces
diff --git a/macq/generate/pddl/vanilla_sampling.py b/macq/generate/pddl/vanilla_sampling.py
index 84d0cbab..2fb550a5 100644
--- a/macq/generate/pddl/vanilla_sampling.py
+++ b/macq/generate/pddl/vanilla_sampling.py
@@ -1,7 +1,13 @@
 from tarski.search.operations import progress
 import random
 from . import Generator
-from ...utils import set_timer_throw_exc, TraceSearchTimeOut, basic_timer, set_num_traces, set_plan_length
+from ...utils import (
+    set_timer_throw_exc,
+    TraceSearchTimeOut,
+    basic_timer,
+    set_num_traces,
+    set_plan_length,
+)
 from ...observation.partial_observation import PercentError
 from ...trace import (
     Step,
@@ -120,5 +126,3 @@ def generate_single_trace(self, plan_len: int = None):
                     trace.append(step)
                     valid_trace = True
         return trace
-
-
diff --git a/macq/observation/identity_observation.py b/macq/observation/identity_observation.py
index f864cf50..65bc5ee5 100644
--- a/macq/observation/identity_observation.py
+++ b/macq/observation/identity_observation.py
@@ -1,6 +1,6 @@
 from dataclasses import dataclass
 from typing import Optional, List
-from ..trace import Step
+from ..trace import Step, State
 from . import Observation, InvalidQueryParameter
 
 
@@ -11,6 +11,8 @@ class IdentityObservation(Observation):
     class.
     """
 
+    state: State
+
     class IdentityState(dict):
         def __hash__(self):
             return hash(tuple(sorted(self.items())))
@@ -42,7 +44,7 @@ def __init__(self, step: Step, **kwargs):
         self.action = None if step.action is None else step.action.clone()
 
     def __hash__(self):
-        return hash(self.details())
+        return super().__hash__()
 
     def __eq__(self, other):
         return (
@@ -51,9 +53,6 @@ def __eq__(self, other):
             and self.action == other.action
         )
 
-    def details(self):
-        return f"Obs {str(self.index)}.\n  State: {str(self.state)}\n  Action: {str(self.action)}"
-
     def _matches(self, key: str, value: str):
         if key == "action":
             if self.action is None:
diff --git a/macq/observation/partial_observation.py b/macq/observation/partial_observation.py
index ca56c132..41e7fa01 100644
--- a/macq/observation/partial_observation.py
+++ b/macq/observation/partial_observation.py
@@ -1,10 +1,9 @@
-from logging import warning
+from warnings import warn
+from typing import Set
 from ..utils import PercentError
 from ..trace import Step, Fluent
 from ..trace import PartialState
 from . import Observation, InvalidQueryParameter
-from typing import Set
-import random
 
 
 class PartialObservation(Observation):
@@ -31,7 +30,7 @@ def __init__(
             raise PercentError()
 
         if percent_missing == 0 and not hide:
-            warning("Creating a PartialObseration with no missing information.")
+            warn("Creating a PartialObseration with no missing information.")
 
         # NOTE: Can't use super due to multiple inheritence (NoisyPartialObservation)
         Observation.__init__(self, index=step.index)
diff --git a/macq/trace/action.py b/macq/trace/action.py
index bdc75cf8..166d58ad 100644
--- a/macq/trace/action.py
+++ b/macq/trace/action.py
@@ -12,8 +12,8 @@ class Action:
     Attributes:
         name (str):
             The name of the action.
-        obj_params (list):
-            The list of objects the action acts on.
+        obj_params (List[PlanningObject]):
+            The set of objects the action acts on.
         cost (int):
             The cost to perform the action.
         precond (Set[Fluent]):
diff --git a/macq/trace/partial_state.py b/macq/trace/partial_state.py
index f4dc18d6..737e8f98 100644
--- a/macq/trace/partial_state.py
+++ b/macq/trace/partial_state.py
@@ -1,16 +1,16 @@
 from . import State
 from . import Fluent
-from typing import Dict
+from typing import Dict, Union
 
 
 class PartialState(State):
     """A Partial State where the value of some fluents are unknown."""
 
-    def __init__(self, fluents: Dict[Fluent, bool] = {}):
+    def __init__(self, fluents: Dict[Fluent, Union[bool, None]] = {}):
         """
         Args:
             fluents (dict):
                 Optional; A mapping of `Fluent` objects to their value in this
                 state. Defaults to an empty `dict`.
         """
-        super().__init__(fluents)
+        self.fluents = fluents
diff --git a/macq/trace/trace.py b/macq/trace/trace.py
index a597126b..314bf1b3 100644
--- a/macq/trace/trace.py
+++ b/macq/trace/trace.py
@@ -54,6 +54,9 @@ def __init__(self, steps: List[Step] = None):
         self.steps = steps if steps is not None else []
         self.__reinit_actions_and_fluents()
 
+    def __eq__(self, other):
+        return isinstance(other, Trace) and self.steps == other.steps
+
     def __len__(self):
         return len(self.steps)
 
@@ -262,7 +265,7 @@ def get_post_states(self, action: Action):
                 post_states.add(self[i + 1].state)
         return post_states
 
-    def get_sas_triples(self, action: Action) -> Set[SAS]:
+    def get_sas_triples(self, action: Action) -> List[SAS]:
         """Retrieves the list of (S,A,S') triples for the action in this trace.
 
         In a (S,A,S') triple, S is the pre-state, A is the action, and S' is
@@ -276,11 +279,11 @@ def get_sas_triples(self, action: Action) -> Set[SAS]:
             A `SAS` object, containing the `pre_state`, `action`, and
             `post_state`.
         """
-        sas_triples = set()
+        sas_triples = []
         for i, step in enumerate(self):
             if step.action == action:
                 triple = SAS(step.state, action, self[i + 1].state)
-                sas_triples.add(triple)
+                sas_triples.append(triple)
         return sas_triples
 
     def get_total_cost(self):
diff --git a/macq/trace/trace_list.py b/macq/trace/trace_list.py
index 2c50722d..e7dd3cd6 100644
--- a/macq/trace/trace_list.py
+++ b/macq/trace/trace_list.py
@@ -1,6 +1,7 @@
+from warnings import warn
+from typing import List, Callable, Type, Optional, Union
 from logging import warn
 from collections.abc import MutableSequence
-from typing import List, Callable, Type, Union
 from rich.console import Console
 
 from . import Action, Trace
diff --git a/macq/utils/__init__.py b/macq/utils/__init__.py
index 84b24c99..1c45496f 100644
--- a/macq/utils/__init__.py
+++ b/macq/utils/__init__.py
@@ -4,6 +4,18 @@
 from .trace_errors import InvalidPlanLength, InvalidNumberOfTraces
 from .trace_utils import set_num_traces, set_plan_length
 from .tokenization_errors import TokenizationError
-#from .tokenization_utils import extract_fluent_subset
 
-__all__ = ["set_timer_throw_exc", "basic_timer", "TraceSearchTimeOut", "ComplexEncoder", "PercentError", "set_num_traces", "set_plan_length", "InvalidPlanLength", "InvalidNumberOfTraces", "TokenizationError",]# "extract_fluent_subset"]
+# from .tokenization_utils import extract_fluent_subset
+
+__all__ = [
+    "set_timer_throw_exc",
+    "basic_timer",
+    "TraceSearchTimeOut",
+    "ComplexEncoder",
+    "PercentError",
+    "set_num_traces",
+    "set_plan_length",
+    "InvalidPlanLength",
+    "InvalidNumberOfTraces",
+    "TokenizationError",
+]
diff --git a/macq/utils/common_errors.py b/macq/utils/common_errors.py
index 681bcc9d..8c8b046f 100644
--- a/macq/utils/common_errors.py
+++ b/macq/utils/common_errors.py
@@ -5,4 +5,4 @@ def __init__(
         self,
         message="The percentage supplied is invalid.",
     ):
-        super().__init__(message)
\ No newline at end of file
+        super().__init__(message)
diff --git a/macq/utils/pysat.py b/macq/utils/pysat.py
new file mode 100644
index 00000000..5a9cf2a5
--- /dev/null
+++ b/macq/utils/pysat.py
@@ -0,0 +1,75 @@
+from typing import List, Tuple, Dict, Hashable
+from pysat.formula import WCNF
+from pysat.examples.rc2 import RC2
+from nnf import And, Or, Var
+
+
+def get_encoding(
+    clauses: And[Or[Var]], start: int = 1
+) -> Tuple[Dict[Hashable, int], Dict[int, Hashable]]:
+    """Maps NNF clauses to pysat clauses and vice-versa.
+
+    Args:
+        clauses (And[Or[Var]]):
+            NNF clauses (in conjunctive normal form) to be mapped to pysat clauses.
+        start (int):
+            Optional; The number to start the mapping from. Defaults to 1.
+
+    Returns:
+        Tuple[Dict[Hashable, int], Dict[int, Hashable]]:
+            The encode mapping (NNF to pysat), and the decode mapping (pysat to NNF).
+    """
+    decode = dict(enumerate(clauses.vars(), start=start))
+    encode = {v: k for k, v in decode.items()}
+    return encode, decode
+
+
+def encode(clauses: And[Or[Var]], encode: Dict[Hashable, int]) -> List[List[int]]:
+    """Encodes NNF clauses into pysat clauses.
+
+    Args:
+        clauses (And[Or[Var]]):
+            NNF clauses (in conjunctive normal form) to be converted to pysat clauses.
+        encode (Dict[Hashable, int]):
+            The encode mapping to apply to the NNF clauses.
+
+    Returns:
+        List[List[int]]:
+            The pysat encoded clauses.
+    """
+    encoded = [
+        [encode[var.name] if var.true else -encode[var.name] for var in clause]
+        for clause in clauses
+    ]
+    return encoded
+
+
+def to_wcnf(
+    soft_clauses: And[Or[Var]], weights: List[int], hard_clauses: And[Or[Var]] = None
+) -> Tuple[WCNF, Dict[int, Hashable]]:
+    """Builds a pysat weighted CNF theory from pysat clauses.
+
+    Args:
+        soft_clauses (And[Or[Var]]):
+            The soft clauses (NNF clauses, in CNF) for the WCNF theory.
+        weights (List[int]):
+            The weights to associate with the soft clauses.
+        hard_clauses (And[Or[Var]]):
+            Optional; Hard clauses (unweighted) to add to the WCNF theory.
+
+    Returns:
+        Tuple[WCNF, Dict[int, Hashable]]:
+            The WCNF theory, and the decode mapping to convert the pysat vars back to NNF.
+    """
+    wcnf = WCNF()
+    soft_encode, decode = get_encoding(soft_clauses)
+    encoded = encode(soft_clauses, soft_encode)
+    wcnf.extend(encoded, weights)
+
+    if hard_clauses:
+        hard_encode, hard_decode = get_encoding(hard_clauses, start=len(decode) + 1)
+        decode.update(hard_decode)
+        encoded = encode(hard_clauses, hard_encode)
+        wcnf.extend(encoded)
+
+    return wcnf, decode
diff --git a/setup.py b/setup.py
index 41619b15..ac4147a0 100644
--- a/setup.py
+++ b/setup.py
@@ -10,6 +10,7 @@
     "tarski@git+git://github.com/aig-upf/tarski.git@devel#egg=tarski[arithmetic]",
     "requests",
     "rich",
+    "nnf",
     "python-sat",
     "bauhaus",
     "numpy"


Test patch:
diff --git a/tests/extract/test_arms.py b/tests/extract/test_arms.py
new file mode 100644
index 00000000..ae9ec82e
--- /dev/null
+++ b/tests/extract/test_arms.py
@@ -0,0 +1,51 @@
+from pathlib import Path
+from typing import List
+from macq.trace import *
+from macq.extract import Extract, modes
+from macq.observation import PartialObservation
+from macq.generate.pddl import *
+
+
+def get_fluent(name: str, objs: List[str]):
+    objects = [PlanningObject(o.split()[0], o.split()[1]) for o in objs]
+    return Fluent(name, objects)
+
+
+def test_arms():
+    base = Path(__file__).parent.parent
+    dom = str((base / "pddl_testing_files/blocks_domain.pddl").resolve())
+    prob = str((base / "pddl_testing_files/blocks_problem.pddl").resolve())
+
+    traces = TraceList()
+    generator = TraceFromGoal(dom=dom, prob=prob)
+
+    generator.change_goal(
+        {
+            get_fluent("on", ["object a", "object b"]),
+            get_fluent("on", ["object b", "object c"]),
+        }
+    )
+    traces.append(generator.generate_trace())
+    generator.change_goal(
+        {
+            get_fluent("on", ["object b", "object a"]),
+            get_fluent("on", ["object c", "object b"]),
+        }
+    )
+    traces.append(generator.generate_trace())
+
+    observations = traces.tokenize(PartialObservation, percent_missing=0.5)
+    model = Extract(
+        observations,
+        modes.ARMS,
+        debug=False,
+        upper_bound=2,
+        min_support=2,
+        action_weight=110,
+        info_weight=100,
+        threshold=0.6,
+        info3_default=30,
+        plan_default=30,
+    )
+
+    assert model
diff --git a/tests/extract/test_extract.py b/tests/extract/test_extract.py
index 1ff02e9f..fa8fc4fe 100644
--- a/tests/extract/test_extract.py
+++ b/tests/extract/test_extract.py
@@ -3,8 +3,7 @@
 from macq.observation import Observation
 from tests.utils.test_traces import blocks_world
 
-# Other functionality of extract is implicitly tested by any extraction technique
-# This is reflected in coverage reports
+# Other functionality of extract is tested by extraction technique tests
 
 
 def test_incompatible_observation_token():
diff --git a/tests/extract/test_observer.py b/tests/extract/test_observer.py
index e00c4aed..2d550a1d 100644
--- a/tests/extract/test_observer.py
+++ b/tests/extract/test_observer.py
@@ -19,12 +19,18 @@ def test_observer():
 if __name__ == "__main__":
     # exit out to the base macq folder so we can get to /tests
     base = Path(__file__).parent.parent
-    model_blocks_dom = str((base / "generated_testing_files/model_blocks_domain.pddl").resolve())
-    model_blocks_prob = str((base / "generated_testing_files/model_blocks_problem.pddl").resolve())
+    model_blocks_dom = str(
+        (base / "generated_testing_files/model_blocks_domain.pddl").resolve()
+    )
+    model_blocks_prob = str(
+        (base / "generated_testing_files/model_blocks_problem.pddl").resolve()
+    )
 
     traces = blocks_world(5)
     observations = traces.tokenize(IdentityObservation)
     traces.print()
     model = Extract(observations, modes.OBSERVER)
     print(model.details())
-    model.to_pddl("model_blocks_dom", "model_blocks_prob", model_blocks_dom, model_blocks_prob)
+    model.to_pddl(
+        "model_blocks_dom", "model_blocks_prob", model_blocks_dom, model_blocks_prob
+    )
diff --git a/tests/extract/test_slaf.py b/tests/extract/test_slaf.py
index 09083f0d..5bbbc08f 100644
--- a/tests/extract/test_slaf.py
+++ b/tests/extract/test_slaf.py
@@ -32,11 +32,10 @@ def test_slaf():
     traces = generate_blocks_traces(plan_len=2, num_traces=1)
     observations = traces.tokenize(
         AtomicPartialObservation,
-        method=AtomicPartialObservation.random_subset,
         percent_missing=0.10,
     )
     traces.print()
-    model = Extract(observations, modes.SLAF, debug_mode=True)
+    model = Extract(observations, modes.SLAF)
     print(model.details())
 
     model.to_pddl(
diff --git a/tests/generate/pddl/test_plan.py b/tests/generate/pddl/test_plan.py
index 0efe17c6..f09e800d 100644
--- a/tests/generate/pddl/test_plan.py
+++ b/tests/generate/pddl/test_plan.py
@@ -20,4 +20,4 @@
     plan = vanilla.generate_plan(from_ipc_file=True, filename=path)
     tracelist = TraceList()
     tracelist.append(vanilla.generate_single_trace_from_plan(plan))
-    tracelist.print(wrap="y")
\ No newline at end of file
+    tracelist.print(wrap="y")
diff --git a/tests/generate/pddl/test_random_goal_sampling.py b/tests/generate/pddl/test_random_goal_sampling.py
index e4b051b1..afd19330 100644
--- a/tests/generate/pddl/test_random_goal_sampling.py
+++ b/tests/generate/pddl/test_random_goal_sampling.py
@@ -14,7 +14,7 @@
         num_traces=3,
         steps_deep=10,
         subset_size_perc=0.1,
-        enforced_hill_climbing_sampling=False
+        enforced_hill_climbing_sampling=False,
     )
     traces = random_sampler.traces
     traces.print(wrap="y")
@@ -27,5 +27,5 @@
         num_traces=3,
         steps_deep=10,
         subset_size_perc=0.1,
-        enforced_hill_climbing_sampling=False
-    ).traces
\ No newline at end of file
+        enforced_hill_climbing_sampling=False,
+    ).traces
diff --git a/tests/generate/pddl/test_trace_from_goal.py b/tests/generate/pddl/test_trace_from_goal.py
index 148a4fb4..25d461c7 100644
--- a/tests/generate/pddl/test_trace_from_goal.py
+++ b/tests/generate/pddl/test_trace_from_goal.py
@@ -32,8 +32,12 @@ def test_invalid_goal_change():
     dom = str((base / "pddl_testing_files/blocks_domain.pddl").resolve())
     prob = str((base / "pddl_testing_files/blocks_problem.pddl").resolve())
 
-    new_blocks_dom = str((base / "generated_testing_files/new_blocks_dom.pddl").resolve())
-    new_blocks_prob = str((base / "generated_testing_files/new_blocks_prob.pddl").resolve())
+    new_blocks_dom = str(
+        (base / "generated_testing_files/new_blocks_dom.pddl").resolve()
+    )
+    new_blocks_prob = str(
+        (base / "generated_testing_files/new_blocks_prob.pddl").resolve()
+    )
     new_game_dom = str((base / "generated_testing_files/new_game_dom.pddl").resolve())
     new_game_prob = str((base / "generated_testing_files/new_game_prob.pddl").resolve())
 
diff --git a/tests/generate/pddl/test_vanilla_sampling.py b/tests/generate/pddl/test_vanilla_sampling.py
index 0765562b..e8deaea8 100644
--- a/tests/generate/pddl/test_vanilla_sampling.py
+++ b/tests/generate/pddl/test_vanilla_sampling.py
@@ -39,8 +39,12 @@ def test_invalid_vanilla_sampling():
     prob = str((base / "pddl_testing_files/blocks_problem.pddl").resolve())
     vanilla = VanillaSampling(dom=dom, prob=prob, plan_len=7, num_traces=10)
 
-    new_blocks_dom = str((base / "generated_testing_files/new_blocks_dom.pddl").resolve())
-    new_blocks_prob = str((base / "generated_testing_files/new_blocks_prob.pddl").resolve())
+    new_blocks_dom = str(
+        (base / "generated_testing_files/new_blocks_dom.pddl").resolve()
+    )
+    new_blocks_prob = str(
+        (base / "generated_testing_files/new_blocks_prob.pddl").resolve()
+    )
     new_game_dom = str((base / "generated_testing_files/new_game_dom.pddl").resolve())
     new_game_prob = str((base / "generated_testing_files/new_game_prob.pddl").resolve())
 
@@ -86,4 +90,6 @@ def test_invalid_vanilla_sampling():
     tracelist.print(wrap="y")
 
     # test generating traces with action preconditions/effects known
-    vanilla_traces = VanillaSampling(problem_id=4627, plan_len=7, num_traces=10, observe_pres_effs=True).traces
+    vanilla_traces = VanillaSampling(
+        problem_id=4627, plan_len=7, num_traces=10, observe_pres_effs=True
+    ).traces
diff --git a/tests/test_readme.py b/tests/test_readme.py
index 217c7305..f81967d7 100644
--- a/tests/test_readme.py
+++ b/tests/test_readme.py
@@ -23,6 +23,7 @@ def test_readme():
 
     action1 = traces[0][0].action
     assert action1
+    assert traces.get_usage(action1)
 
     trace = traces[0]
     assert len(trace) == 5
@@ -44,4 +45,3 @@ def test_readme():
     traces = generate_traces()
     model = extract_model(traces)
     print(model.details())
-
diff --git a/tests/trace/test_trace.py b/tests/trace/test_trace.py
index 6e4c9be3..c21eb43d 100644
--- a/tests/trace/test_trace.py
+++ b/tests/trace/test_trace.py
@@ -61,7 +61,7 @@ def test_trace_get_sas_triples():
     (state2, state3) = (trace.steps[1].state, trace.steps[2].state)
 
     assert isinstance(action2, Action)
-    assert trace.get_sas_triples(action2) == {SAS(state2, action2, state3)}
+    assert trace.get_sas_triples(action2) == [SAS(state2, action2, state3)]
 
 
 # test that the total cost is working correctly