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

Issue:
OBSERVER (AIPS'94)
## Notes
- The original approach

## References
- [x] [Learning Planning Operators by Observation and Practice](https://aaai.org/Papers/AIPS/1994/AIPS94-057.pdf)
Fix tests
<!-- Change ## at the end of the url to the pull request number -->
![Coverage Badge](https://img.shields.io/endpoint?url=https://gist.githubusercontent.com/haz/03ac305b42d7c9ad4ef3213341bf3f2f/raw/macq__pull_78.json?cacheSeconds=3500)



Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index 3cc653d1..0e0a18b2 100644
--- a/README.md
+++ b/README.md
@@ -17,8 +17,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 import generate
 
 # get a domain-specific generator: uses api.planning.domains problem_id/
 # generate 100 traces of length 20 using vanilla sampling
@@ -34,20 +33,36 @@ trace.actions
 trace.get_prev_states(action) # get the state before each occurance of action
 trace.get_post_states(action) # state after each occurance of action
 trace.get_total_cost()
-trace.tokenize(IdentityObservation)
-trace.tokens
 
 step = trace[0]
 step.base_fluents()
 step.base_action()
 
-model = extract.Extract(traces, extract.modes.OBSERVER)
-
-model.actions
+######################################################################
+# Model Extraction
+######################################################################
+from macq import extract
+from macq.observation import IdentityObservation
+observations = traces.tokenize(IdentityObservation)
+model = extract.Extract(observations, extract.modes.OBSERVER)
+model.details()
+
+Model:
+  Fluents: on table block A, holding block B, on table block B, clear block A, holding block A, on block A block B, on block A block B, clear block B
+  Actions:
+    pick up block A:
+      precond:
+        on table block B
+        clear block A
+      add:
+        holding block A
+        clear block B
+      delete:
+        on table block A
+        on block A block B
+    stack block B block A:
 ...
 
-model.actions[0].to_pddl()
-
 ```
 
 ## Coverage <a name="coverage"></a>
diff --git a/macq/extract/__init__.py b/macq/extract/__init__.py
index a9a2c5b3..07d5c0f2 100644
--- a/macq/extract/__init__.py
+++ b/macq/extract/__init__.py
@@ -1,1 +1,10 @@
-__all__ = []
+from .model import Model, ModelAction
+from .extract import Extract, modes, IncompatibleObservationToken
+
+__all__ = [
+    "Model",
+    "Extract",
+    "modes",
+    "IncompatibleObservationToken",
+    "ModelAction",
+]
diff --git a/macq/extract/extract.py b/macq/extract/extract.py
new file mode 100644
index 00000000..e01cbd27
--- /dev/null
+++ b/macq/extract/extract.py
@@ -0,0 +1,60 @@
+from dataclasses import dataclass
+from typing import List
+from enum import Enum, auto
+from .observer import Observer
+from ..observation import Observation
+from ..trace import ObservationList, Action, State
+
+
+@dataclass
+class SAS:
+    pre_state: State
+    action: Action
+    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.
+
+    An Enum where each value represents a model extraction technique.
+    """
+
+    OBSERVER = auto()
+
+
+class Extract:
+    """Extracts models from observations.
+
+    The Extract class uses an extraction method to retrieve an action Model
+    from state observations.
+    """
+
+    def __new__(cls, observations: ObservationList, mode: modes):
+        """Extracts a Model object.
+
+        Extracts a model from the observations using the specified extraction
+        technique.
+
+        Args:
+            observations (ObservationList):
+                The state observations to extract the model from.
+            mode (Enum):
+                The extraction technique to use.
+
+        Returns:
+            A Model object. The model's characteristics are determined by the
+            extraction technique used.
+        """
+
+        techniques = {
+            modes.OBSERVER: Observer,
+        }
+
+        return techniques[mode](observations)
diff --git a/macq/extract/model.py b/macq/extract/model.py
new file mode 100644
index 00000000..faa8216f
--- /dev/null
+++ b/macq/extract/model.py
@@ -0,0 +1,99 @@
+from json import dumps, loads
+from .model_action import ModelAction
+from ..trace import Fluent
+from typing import Set
+
+
+class Model:
+    """Action model.
+
+    An action model representing a planning domain. The characteristics of the
+    model are dependent on the extraction technique used to obtain the model.
+
+    Attributes:
+        fluents (set):
+            The set of fluents in the domain.
+        actions (set):
+            The set of actions in the domain. Actions include their
+            preconditions, add effects, and delete effects. The nature of the
+            action attributes characterize the model.
+    """
+
+    def __init__(self, fluents: Set[Fluent], actions: Set[ModelAction]):
+        """Initializes a Model with a set of fluents and a set of actions.
+
+        Args:
+            fluents (set):
+                The set of fluents in the model.
+            actions (set):
+                The set of actions in the model.
+        """
+        self.fluents = fluents
+        self.actions = actions
+
+    def details(self):
+        # Set the indent width
+        indent = " " * 2
+        string = "Model:\n"
+        # Map fluents to a comma separated string of the fluent names
+        string += f"{indent}Fluents: {', '.join([f.details() for f in self.fluents])}\n"
+        # Map the actions to a summary of their names, preconditions, add
+        # effects and delete effects
+        string += f"{indent}Actions:\n"
+        for line in self._get_action_details().splitlines():
+            string += f"{indent * 2}{line}\n"
+        return string
+
+    def _get_action_details(self):
+        # Set the indent width
+        indent = " " * 2
+        details = ""
+        for action in self.actions:
+            details += action.details() + ":\n"
+            details += f"{indent}precond:\n"
+            for f in action.precond:
+                details += f"{indent * 2}{f.details()}\n"
+            details += f"{indent}add:\n"
+            for f in action.add:
+                details += f"{indent * 2}{f.details()}\n"
+            details += f"{indent}delete:\n"
+            for f in action.delete:
+                details += f"{indent * 2}{f.details()}\n"
+
+        return details
+
+    def serialize(self, filepath: str = None):
+        """Serializes the model into a json string.
+
+        Args:
+            filepath (str):
+                Optional; If supplied, the json string will be written to the
+                filepath.
+
+        Returns:
+            A string in json format representing the model.
+        """
+        serial = dumps(self, indent=2, default=lambda o: o.__dict__)
+        if filepath is not None:
+            with open(filepath, "w") as fp:
+                fp.write(serial)
+        return serial
+
+    @staticmethod
+    def deserialize(string: str):
+        """Deserializes a json string into a Model.
+
+        Args:
+            string (str):
+                The json string representing a model.
+
+        Returns:
+            A Model object matching the one specified by `string`.
+        """
+        return Model._from_json(loads(string))
+
+    @classmethod
+    def _from_json(cls, data: dict):
+        fluents = set(map(Fluent.from_json, data["fluents"]))
+        actions = set(map(ModelAction.from_json, data["actions"]))
+        return cls(fluents, actions)
diff --git a/macq/extract/model_action.py b/macq/extract/model_action.py
new file mode 100644
index 00000000..55508c82
--- /dev/null
+++ b/macq/extract/model_action.py
@@ -0,0 +1,61 @@
+from typing import Set
+from ..trace import Action, Fluent, PlanningObject
+
+
+class ModelAction:
+    def __init__(self, action: Action):
+        self.name = action.name
+        self.obj_params = action.obj_params
+        self.cost = action.cost
+        self.precond = set()
+        self.add = set()
+        self.delete = set()
+
+    def __eq__(self, other):
+        if not isinstance(other, ModelAction):
+            return False
+        return self.name == other.name and self.obj_params == other.obj_params
+
+    def __hash__(self):
+        # Order of obj_params is important!
+        return hash(self.details())
+
+    def details(self):
+        string = f"{self.name} {' '.join([o.details() for o in self.obj_params])}"
+        return string
+
+    def update_precond(self, fluents: Set[Fluent]):
+        """Adds preconditions to the action.
+
+        Args:
+            fluents (set):
+                The set of fluents to be added to the action's preconditions.
+        """
+        self.precond.update(fluents)
+
+    def update_add(self, fluents: Set[Fluent]):
+        """Adds add effects to the action.
+
+        Args:
+            fluents (set):
+                The set of fluents to be added to the action's add effects.
+        """
+        self.add.update(fluents)
+
+    def update_delete(self, fluents: Set[Fluent]):
+        """Adds delete effects to the action.
+
+        Args:
+            fluents (set):
+                The set of fluents to be added to the action's delete effects.
+        """
+        self.delete.update(fluents)
+
+    @classmethod
+    def from_json(cls, data):
+        """Converts a json object to an Action."""
+        obj_params = list(map(PlanningObject.from_json, data["obj_params"]))
+        precond = set(map(Fluent.from_json, data["precond"]))
+        add = set(map(Fluent.from_json, data["add"]))
+        delete = set(map(Fluent.from_json, data["delete"]))
+        return cls(data["name"], obj_params, data["cost"], precond, add, delete)
diff --git a/macq/extract/observer.py b/macq/extract/observer.py
new file mode 100644
index 00000000..9c9f1cbd
--- /dev/null
+++ b/macq/extract/observer.py
@@ -0,0 +1,90 @@
+from typing import List, Set
+from collections import defaultdict
+import macq.extract as extract
+from .model import Model
+from ..trace import ObservationList, Action
+from ..observation import Observation, IdentityObservation
+
+
+class Observer:
+    """Observer model extraction method.
+
+    Extracts a Model from state observations using the observer technique.  All
+    the fluents present in the observations are saved as is. The actions have
+    their preconditions, add effects, and delete effects found and added to the
+    object.
+
+    The preconditions of an action are defined as the (positive) intersection
+    of all pre-states to that action. The effects are the union off all fluents
+    that changed state from a pre-state to a post-state of an action. Add this
+    effects are fluents that went from False to True, delete effects are
+    fluents that went from True to False.
+    """
+
+    def __new__(cls, observations: ObservationList):
+        """Creates a new Model object.
+
+        Args:
+            observations (ObservationList):
+                The state observations to extract the model from.
+        Raises:
+            IncompatibleObservationToken:
+                Raised if the observations are not identity observation.
+        """
+        if observations.type is not IdentityObservation:
+            raise extract.IncompatibleObservationToken(observations.type, Observer)
+        fluents = Observer._get_fluents(observations)
+        actions = Observer._get_actions(observations)
+        return Model(fluents, actions)
+
+    @staticmethod
+    def _get_fluents(observations: ObservationList):
+        """Retrieves the set of fluents in the observations."""
+        fluents = set()
+        trace: List[IdentityObservation]
+        for trace in observations:
+            for obs in trace:
+                # Update fluents with the fluents in this observation
+                fluents.update(list(obs.state.keys()))
+        return fluents
+
+    @staticmethod
+    def _get_actions(observations: ObservationList):
+        """Retrieves and augments the set of actions in the observations."""
+        # Get the unique actions and the relevant traces
+        action_obs = defaultdict(list)
+        trace_obs: List[IdentityObservation]
+        for trace_obs in observations:
+            for obs in trace_obs:
+                action = obs.action
+                if action is not None:  # Final step has no action
+                    action_obs[action].append(trace_obs)
+
+        # Create the ModelActions
+        action_pre_states = defaultdict(set)
+        action_transitions = observations.get_all_transitions()
+        for action, transitions in action_transitions.items():
+            model_action = extract.ModelAction(action)
+            for pre, post in transitions:
+                # Add all action pre-states to a set
+                action_pre_states[model_action].add(pre.state)
+                # Update the action's effects
+                delta = pre.state.diff_from(post.state)
+                model_action.update_add(delta.added)
+                model_action.update_delete(delta.deleted)
+
+        for action, pre_states in action_pre_states.items():
+            # Find the (positive) intersection of the pre-states
+            precond = set.intersection(*map(Observer._filter_positive, pre_states))
+            action.update_precond(precond)
+
+        return {action for action in action_pre_states.keys()}
+
+    @staticmethod
+    def _filter_positive(state):
+        """Returns the set of true fluents in a state."""
+        true_fluents = set()
+        for fluent, is_true in state.items():
+            if is_true:
+                true_fluents.add(fluent)
+        return true_fluents
diff --git a/macq/generate/pddl/__init__.py b/macq/generate/pddl/__init__.py
index 72dbd023..b1926d6b 100644
--- a/macq/generate/pddl/__init__.py
+++ b/macq/generate/pddl/__init__.py
@@ -1,4 +1,4 @@
 from .generator import Generator
 from .vanilla_sampling import VanillaSampling
 
-__all__ = ["Generator", "VanillaSampling"]
\ No newline at end of file
+__all__ = ["Generator", "VanillaSampling"]
diff --git a/macq/generate/pddl/generator.py b/macq/generate/pddl/generator.py
index f0d5e1af..046be9fa 100644
--- a/macq/generate/pddl/generator.py
+++ b/macq/generate/pddl/generator.py
@@ -1,4 +1,4 @@
-from ...trace import Action, State, CustomObject, Fluent
+from ...trace import Action, State, PlanningObject, Fluent
 from tarski.io import PDDLReader, FstripsWriter
 from tarski.search import GroundForwardSearchModel
 from tarski.grounding.lp_grounding import ground_problem_schemas_into_plain_operators
@@ -35,9 +35,9 @@ def __init__(self, dom: str = "", prob: str = "", problem_id: int = None):
         self.instance = GroundForwardSearchModel(self.problem, operators)
 
         """
-        Class that handles creating a basic PDDL state trace generator. Handles all 
+        Class that handles creating a basic PDDL state trace generator. Handles all
         parsing and stores the problem, language, and grounded instance for the child
-        generators to easily access and use. Takes either the raw filenames of the 
+        generators to easily access and use. Takes either the raw filenames of the
         domain and problem, or a problem ID.
 
         Arguments
@@ -145,7 +145,7 @@ def __tarski_atom_to_macq_fluent(self, atom: Atom):
         for term in terms:
             if isinstance(fluent_name, BuiltinPredicateSymbol):
                 fluent_name = fluent_name.value
-            objects.append(CustomObject(term.sort.name, term.name))
+            objects.append(PlanningObject(term.sort.name, term.name))
         fluent = Fluent(fluent_name, objects)
         return fluent
 
@@ -207,4 +207,4 @@ def tarski_act_to_macq(self, tarski_act: PlainOperator):
             objs.update(set(fluent.objects))
         for fluent in precond:
             objs.update(set(fluent.objects))
-        return Action(name, objs, precond, add, delete)
+        return Action(name, objs)
diff --git a/macq/observation/__init__.py b/macq/observation/__init__.py
index 13fa34ce..bde0e883 100644
--- a/macq/observation/__init__.py
+++ b/macq/observation/__init__.py
@@ -1,4 +1,4 @@
-from .observation import Observation
+from .observation import Observation, InvalidQueryParameter
 from .identity_observation import IdentityObservation
 
-__all__ = ["Observation", "IdentityObservation"]
+__all__ = ["Observation", "IdentityObservation", "InvalidQueryParameter"]
diff --git a/macq/observation/identity_observation.py b/macq/observation/identity_observation.py
index fcee29b3..1a85650a 100644
--- a/macq/observation/identity_observation.py
+++ b/macq/observation/identity_observation.py
@@ -1,5 +1,5 @@
 from ..trace import Step
-from . import Observation
+from . import Observation, InvalidQueryParameter
 
 
 class IdentityObservation(Observation):
@@ -7,7 +7,7 @@ class IdentityObservation(Observation):
     The identity observation stores the step unmodified.
     """
 
-    def __init__(self, step: Step):
+    def __init__(self, step: Step, **kwargs):
         """
         Creates an IdentityObservation object, storing the step.
 
@@ -16,9 +16,27 @@ def __init__(self, step: Step):
         step : Step
             The step associated with this observation.
         """
-        super().__init__(step)
+        super().__init__(index=step.index, **kwargs)
+        self.state = step.state.clone()
+        self.action = None if step.action is None else step.action.clone()
 
-    def __eq__(self, value):
-        if isinstance(value, IdentityObservation):
-            return self.step == value.step
-        return False
+    def __hash__(self):
+        return hash(self.details())
+
+    def __eq__(self, other):
+        if not isinstance(other, IdentityObservation):
+            return False
+        return self.state == other.state and self.action == other.action
+
+    def _matches(self, key: str, value: str):
+        if key == "action":
+            if self.action is None:
+                return value is None
+            return self.action.details() == value
+        elif key == "fluent_holds":
+            return self.state.holds(value)  # whatever this needs to look like
+        else:
+            raise InvalidQueryParameter(IdentityObservation, key)
+
+    def details(self):
+        return str(self.index) + self.action.details() + self.state.details()
diff --git a/macq/observation/observation.py b/macq/observation/observation.py
index efa0332e..179a9647 100644
--- a/macq/observation/observation.py
+++ b/macq/observation/observation.py
@@ -1,15 +1,31 @@
+from logging import warn
 from ..trace import Step
 
 
+class InvalidQueryParameter(Exception):
+    def __init__(self, obs, param, message=None):
+        if message is None:
+            message = f"{param} is not a valid query parameter for {obs.__name__}"
+        super().__init__(message)
+
+
 class Observation:
     """
     An Observation object stores an observation token representation of a step.
     """
 
-    def __init__(self, step: Step):
+    def __init__(self, **kwargs):
         """
         Creates an Observation object, storing the step as a token, as well as its index/"place"
         in the trace (which corresponds to that of the step).
         """
-        self.step = step
-        self.index = step.index
+        if "index" in kwargs.keys():
+            self.index = kwargs["index"]
+        else:
+            warn("Creating an Observation token without an index.")
+
+    def _matches(self, *_):
+        raise NotImplementedError()
+
+    def matches(self, query: dict):
+        return all([self._matches(key, value) for key, value in query.items()])
diff --git a/macq/trace/__init__.py b/macq/trace/__init__.py
index e6947717..dd1a0727 100644
--- a/macq/trace/__init__.py
+++ b/macq/trace/__init__.py
@@ -1,9 +1,22 @@
-from .action import Action, Fluent, CustomObject
-from .state import State
-from .step import Step
-from .trace import Trace
-from .trace_list import TraceList
+from .action import Action, PlanningObject
+from .fluent import Fluent
+from .state import State, DeltaState
 from .partially_observable_state import PartiallyObservableState
+from .step import Step
+from .trace import Trace, SAS
+from .trace_list import TraceList, ObservationList
 
 
-__all__ = ["Action", "Fluent", "CustomObject", "State", "Step", "Trace", "TraceList", "PartiallyObservableState"]
+__all__ = [
+    "Action",
+    "Fluent",
+    "PlanningObject",
+    "State",
+    "PartiallyObservableState",
+    "DeltaState",
+    "Step",
+    "Trace",
+    "SAS",
+    "TraceList",
+    "ObservationList",
+]
diff --git a/macq/trace/action.py b/macq/trace/action.py
index 96ffd8b7..1c208550 100644
--- a/macq/trace/action.py
+++ b/macq/trace/action.py
@@ -1,147 +1,62 @@
 from typing import List
-from .fluent import CustomObject, Fluent
+from .fluent import PlanningObject
 
 
 class Action:
-    class InvalidFluent(Exception):
-        """
-        Raised when the user attempts to add fluents (to a precondition or effect) that act on objects
-        outside of the parameters supplied to the action.
-        """
-
-        def __init__(
-            self,
-            message="Cannot add a fluent referencing objects beyond the scope of this action.",
-        ):
-            super().__init__(message)
+    """Grounded action.
+    An Action represents a grounded action in a Trace or a Model. The action's
+    `precond`, `add`, and `delete` attributes characterize a Model, and are
+    found during model extraction.
+    Attributes:
+        name (str):
+            The name of the action.
+        obj_params (list):
+            The list of objects the action acts on.
+        cost (int):
+            The cost to perform the action.
+    """
 
     def __init__(
         self,
         name: str,
-        obj_params: List[CustomObject],
-        precond: List[Fluent],
-        add: List[Fluent],
-        delete: List[Fluent],
+        obj_params: List[PlanningObject],
         cost: int = 0,
     ):
-        """
-        Class to handle each action.
-
-        Arguments
-        ---------
-        name : str
-            The name of the action.
-        obj_params : list
-            The list of objects this action can apply to.
-        precond: list of Fluents
-            The list of preconditions that must pass before this action takes place.
-        add : list of Fluents
-            The list of fluents to be flipped to True after this action takes place.
-        delete : list of Fluents
-            The list of fluents to be flipped to False after this action takes place.
-        cost : int
-            The cost of this action.
+        """Initializes an Action with the parameters provided.
+        The `precond`, `add`, and `delete` args should only be provided in
+        Model deserialization.
+        Args:
+            name (str):
+                The name of the action.
+            obj_params (list):
+                The list of objects the action acts on.
+            cost (int):
+                Optional; The cost to perform the action. Defaults to 0.
         """
         self.name = name
         self.obj_params = obj_params
-        self.precond = []
-        self.add_precond(precond)
-        self.add = []
-        self.add_effect_add(add)
-        self.delete = []
-        self.add_effect_delete(delete)
         self.cost = cost
 
-    def __repr__(self):
-        string = "Action with Name: " + self.name + "\n\nObject Parameters:\n"
-        for obj in self.obj_params:
-            string += str(obj) + "\n"
-        string += "\nPreconditions:\n"
-        for fluent in self.precond:
-            string += str(fluent) + "\n\n"
-        string += "\nEffects to add:\n"
-        for fluent in self.add:
-            string += str(fluent) + "\n\n"
-        string += "\nEffects to delete:\n"
-        for fluent in self.delete:
-            string += str(fluent) + "\n\n"
-        string += "\nCost: " + str(self.cost)
-        return string
-
-    def __add_fluent(self, fluents: List[Fluent], condition: List[Fluent]):
-        """
-        Checks the validity of a fluent before adding it to either the action's preconditions,
-        add effects or delete effects.
-
-        Arguments
-        ---------
-        fluents : list of Fluents
-            The list of fluents to be added to the given action condition.
-        condition : list of Fluents
-            Either the action's preconditions, add effects, or delete effects to be added to.
-        """
-        valid = False
-        for fluent in fluents:
-            for obj in fluent.objects:
-                # if this fluent takes no parameters it is automatically valid
-                if (
-                    len(fluent.objects) == 1
-                    and fluent.objects[0].name == ""
-                    and fluent.objects[0].obj_type == ""
-                ):
-                    valid = True
-                else:
-                    for param in self.obj_params:
-                        if obj == param:
-                            valid = True
-                            # no need to check the rest
-                            break
-                if not valid:
-                    raise self.InvalidFluent()
-                # reset flag
-                valid = False
-        condition.extend(fluents)
-
-    def add_precond(self, fluents: List[Fluent]):
-        """
-        Adds the specified list of fluents to the action's preconditions.
-
-        Arguments
-        ---------
-        fluents : list of Fluents
-            The list of fluents to be added to the action's preconditions.
-        """
-        self.__add_fluent(fluents, self.precond)
-
-    def add_effect_add(self, fluents: List[Fluent]):
-        """
-        Adds the specified list of fluents to the action's add effects.
-
-        Arguments
-        ---------
-        fluents : list of Fluents
-            The list of fluents to be added to the action's add effects.
-        """
-        self.__add_fluent(fluents, self.add)
+    def __hash__(self):
+        # Order of obj_params is important!
+        return hash(self.details())
 
-    def add_effect_delete(self, fluents: List[Fluent]):
-        """
-        Adds the specified list of fluents to the action's delete effects.
+    def __eq__(self, other):
+        if not isinstance(other, Action):
+            return False
+        return self.name == other.name and self.obj_params == other.obj_params
 
-        Arguments
-        ---------
-        fluents : list of Fluents
-            The list of fluents to be added to the action's delete effects.
-        """
-        self.__add_fluent(fluents, self.delete)
+    def details(self):
+        string = f"{self.name} {' '.join([o.details() for o in self.obj_params])}"
+        return string
 
-    def add_parameter(self, obj: CustomObject):
-        """
-        Adds the specified object to the action's list of available parameters.
+    def clone(self):
+        return Action(self.name, self.obj_params, self.cost)
 
-        Arguments
-        ---------
-        obj : CustomObject
-            The object to be added to the action's list of available parameters.
+    def add_parameter(self, obj: PlanningObject):
+        """Adds an object to the action's parameters.
+        Args:
+            obj (PlanningObject):
+                The object to be added to the action's object parameters.
         """
         self.obj_params.append(obj)
diff --git a/macq/trace/fluent.py b/macq/trace/fluent.py
index e9e0d1a1..22fe1dfc 100644
--- a/macq/trace/fluent.py
+++ b/macq/trace/fluent.py
@@ -1,43 +1,78 @@
 from typing import List
 
 
-class CustomObject:
+class PlanningObject:
+    """Objects of a planning domain.
+
+    Attributes:
+        obj_type (str):
+            The type of object in the problem domain.
+            Example: "block".
+        name (str):
+            The name of the object.
+            Example: "A"
+    """
+
     def __init__(self, obj_type: str, name: str):
+        """Initializes a PlanningObject with a type and a name.
+
+        Args:
+            obj_type (str):
+                The type of object in the problem domain.
+            name (str):
+                The name of the object.
+        """
         self.obj_type = obj_type
         self.name = name
 
-    def __repr__(self):
-        string = "Type: " + self.obj_type + ", Name: " + self.name
-        return string
+    def __hash__(self):
+        return hash(self.name)
 
     def __eq__(self, other):
-        return isinstance(other, CustomObject) and self.name == other.name
+        return isinstance(other, PlanningObject) and self.name == other.name
 
-    def __hash__(self):
-        return hash(self.name)
+    def details(self):
+        return " ".join([self.obj_type, self.name])
+
+    @classmethod
+    def from_json(cls, data):
+        """Converts a json object to a PlanningObject."""
+        return cls(**data)
 
 
 class Fluent:
-    def __init__(self, name: str, objects: List[CustomObject], value: bool = True):
-        """
-        Class to handle a predicate and the objects it is applied to.
-
-        Arguments
-        ---------
-        name : str
-            The name of the predicate.
-        objects : list
-            The list of objects this predicate applies to.
-        value : bool
-            The value of the fluent (true or false)
+    """Fluents of a planning domain.
+
+    Attributes:
+        name (str):
+            The name of the fluent.
+            Example: "holding".
+        objects (list):
+            The objects this fluent applies to.
+            Example: Block A.
+    """
+
+    def __init__(self, name: str, objects: List[PlanningObject]):
+        """Initializes a Fluent with a name and a list of objects.
+
+        Args:
+        name (str):
+            The name of the fluent.
+        objects (list):
+            The objects this fluent applies to.
         """
         self.name = name
         self.objects = objects
-        self.value = value
-
-    def __repr__(self):
-        string = "Fluent with Name: " + self.name + "\nObjects:\n"
-        for obj in self.objects:
-            string += str(obj) + "\n"
-        string += "Value: " + str(self.value)
-        return string
+
+    def __hash__(self):
+        # Order of objects is important!
+        return hash(self.details())
+
+    def details(self):
+        return f"{self.name} {' '.join([o.details() for o in self.objects])}"
+
+    @classmethod
+    def from_json(cls, data):
+        """Converts a json object to a Fluent."""
+        objects = list(map(PlanningObject.from_json, data["objects"]))
+        return cls(data["name"], objects)
diff --git a/macq/trace/state.py b/macq/trace/state.py
index eb7c1026..585b3f29 100644
--- a/macq/trace/state.py
+++ b/macq/trace/state.py
@@ -1,23 +1,116 @@
-from typing import List
+from __future__ import annotations
+from dataclasses import dataclass
 from . import Fluent
 
 
+@dataclass
+class DeltaState:
+    added: set[Fluent]
+    deleted: set[Fluent]
+
+
 class State:
-    """
-    Class for a State, which is the set of all fluents and their values at a particular Step.
+    """State in a trace.
+
+    A dict-like object. Maps `Fluent` objects to boolean values, representing
+    the state for a `Step` in a `Trace`.
 
-    Arguments
-    ---------
-    fluents : List of Fluents
-            A list of fluents representing the state.
+    Attributes:
+        fluents (dict):
+            A mapping of `Fluent` objects to their value in this state.
     """
 
-    def __init__(self, fluents: List[Fluent]):
+    def __init__(self, fluents: dict[Fluent, bool] = {}):
+        """Initializes State with an optional fluent-value mapping.
+
+        Args:
+            fluents (dict):
+                Optional; A mapping of `Fluent` objects to their value in this
+                state. Defaults to an empty `dict`.
+        """
         self.fluents = fluents
 
-    def __repr__(self):
-        string = "State:\n\n"
-        for fluent in self.fluents:
-            string += str(fluent) + "\n"
-        string = string.strip()
-        return string
+    def __eq__(self, other):
+        if not isinstance(other, State):
+            return False
+        return self.fluents == other.fluents
+
+    def __hash__(self):
+        return hash(self.details())
+
+    def __len__(self):
+        return len(self.fluents)
+
+    def __setitem__(self, key: Fluent, value: bool):
+        self.fluents[key] = value
+
+    def __getitem__(self, key: Fluent):
+        return self.fluents[key]
+
+    def __delitem__(self, key: Fluent):
+        del self.fluents[key]
+
+    def __iter__(self):
+        return iter(self.fluents)
+
+    def __contains__(self, key):
+        return self.fluents[key]
+
+    def clear(self):
+        return self.fluents.clear()
+
+    def copy(self):
+        return self.fluents.copy()
+
+    def has_key(self, k):
+        return k in self.fluents
+
+    def update(self, *args, **kwargs):
+        return self.fluents.update(*args, **kwargs)
+
+    def keys(self):
+        return self.fluents.keys()
+
+    def values(self):
+        return self.fluents.values()
+
+    def items(self):
+        return self.fluents.items()
+
+    def details(self):
+        string = ""
+        for fluent, value in self.items():
+            string += f"{fluent.details()} ({value}), "
+        return string[:-2]
+
+    def clone(self):
+        return State(self.fluents)
+
+    def holds(self, fluent: str):
+        fluents = dict(map(lambda f: (f.name, f), self.keys()))
+        if fluent in fluents.keys():
+            return self[fluents[fluent]]
+
+    def diff_from(self, other: State):
+        """Determines the delta-state between this state and `other`.
+
+        Args:
+            other (State):
+                The state to compare this one to.
+
+        Returns:
+            A `DeltaState` object, containing two sets: `added` and `deleted`.
+            The added set contains the list of fluents that were False in this
+            state and True in `other`. The deleted set contains the list of
+            fluents that were True in this state and False in `other`.
+        """
+        added = set()
+        deleted = set()
+        fluents = list(self.keys())
+        fluents.extend(list(other.keys()))
+        for f in fluents:
+            if self[f] and not other[f]:  # true pre, false post -> deleted
+                deleted.add(f)
+            elif not self[f] and other[f]:  # false pre, true post -> added
+                added.add(f)
+        return DeltaState(added, deleted)
diff --git a/macq/trace/step.py b/macq/trace/step.py
index 696d6ced..93b316c7 100644
--- a/macq/trace/step.py
+++ b/macq/trace/step.py
@@ -1,56 +1,37 @@
+from typing import Optional
 from . import Action
 from . import State
 
 
 class Step:
-    """
-    A Step object stores the action, and state prior to the action for a step
-    in a trace.
-    """
+    """A step in a Trace.
 
-    def __init__(self, action: Action, state: State, index: int):
-        """
-        Creates a Step object. This stores action, and state prior to the
-        action.
+    A Step object stores a State and the action that is taken from the state.
+    The final step in a trace will not have an action associated with it.
+
+    Attributes:
+        state (State):
+            The state in this step.
+        action (Action | None):
+            The action taken from the state in this step.
+    """
 
-        Attributes
-        ----------
-        action : Action
-            The action taken in this step.
-        state : State
-            The state (list of fluents) at this step.
-        index : int
-            The placement of this step within the trace.
+    def __init__(self, state: State, action: Optional[Action], index: int):
+        """Initializes a Step with a state and optionally an action.
+
+        Args:
+            state (State):
+                The state in this step.
+            action (Action | None):
+                The action taken from the state in this step. Must provide a
+                value, but value can be None.
+            index (int):
+                The index of this step in the trace.
         """
-        self.action = action
         self.state = state
+        self.action = action
         self.index = index
 
-    def __repr__(self):
-        string = str(self.action) + "\n\n" + str(self.state)
+    def details(self):
+        string = self.action.details() + "\n\n" + self.state.details()
         return string
-
-    def base_fluents(self):
-        """
-        Retrieves the names of all fluents used in this step.
-
-        Returns
-        -------
-        list : str
-            Returns a list of the names of all fluents used in this step.
-        """
-        fluents = []
-        for fluent in self.state.fluents:
-            fluents.append(fluent.name)
-        return fluents
-
-    def base_action(self):
-        """
-        Retrieves the name of the action used in this step.
-
-        Returns
-        -------
-        list : str
-            Returns the name of the action used in this step.
-        """
-        return self.action.name
diff --git a/macq/trace/trace.py b/macq/trace/trace.py
index bc445ba9..8ad8d0c9 100644
--- a/macq/trace/trace.py
+++ b/macq/trace/trace.py
@@ -1,62 +1,77 @@
-from typing import List, Type, Iterable, Callable
-from . import Action
-from . import Step
-from . import State
+from dataclasses import dataclass
+from typing import List, Type, Iterable, Callable, Set
+from inspect import cleandoc
+from . import Action, Step, State
 from ..observation import Observation
 
 
+@dataclass
+class SAS:
+    pre_state: State
+    action: Action
+    post_state: State
+
+    def __hash__(self):
+        return hash(
+            self.pre_state.details() + self.action.details() + self.post_state.details()
+        )
+
+
 class Trace:
-    """
-    Class for a Trace, which consists of each Step in a generated solution.
+    """A state trace of a planning problem.
 
+    A `list`-like object, where each element is a step of the state trace.
+
+    Attributes:
+        steps (list):
+            The list of Step objcts constituting the trace.
+        fluents (set):
+            The set of fluents in the trace.
+        actions (set):
+            The set of actions in the trace.
     """
 
     class InvalidCostRange(Exception):
-        """
-        Exception raised for incorrect user input for the cost range.
-        """
-
         def __init__(self, message):
             super().__init__(message)
 
-    def __init__(self, steps: List[Step] = []):
-        """
-        Creates a Trace object.
-
-        Arguments
-        ---------
-        steps : List of Steps (optional)
-            The list of Step objects that make up the trace.
-
-        Attributes
-        ----------
-        fluents : List of str
-            The list of the names of all fluents used.
-            Information on the values of fluents are found in the steps.
-        actions: List of Actions
-            The list of the names of all actions used.
-            Information on the preconditions/effects of actions are found in the steps.
-        observations: List of Observations
-            The set of observation tokens, tokenized from the steps.
+    def __init__(self, steps: List[Step] = None):
+        """Initializes a Trace with an optional list of steps.
+
+        Args:
+            steps (list):
+                Optional; The list of steps in the trace. Defaults to an empty
+                `list`.
         """
-        self.steps = steps
+        self.steps = steps if steps is not None else []
         self.__reinit_actions_and_fluents()
-        self.observations = []
-
-    def __repr__(self):
-        string = "TRACE:\n\nAttributes:"
-        string += "\n\nBase Fluents:\n"
-        for fluent in self.fluents:
-            string += fluent + "\n"
-        string += "\nBase Actions:\n"
-        for action in self.actions:
-            string += action + "\n"
-        string += "\n\nSteps:"
-        for i in range(len(self.steps)):
-            string += "\n\nSTEP " + str(i + 1) + ":\n\n" + str(self.steps[i]) + "\n"
-        string += "\nObservations:\n"
-        for obs in self.observations:
-            string += "\n" + str(obs)
+
+    def details(self):
+        indent = " " * 2
+        # Summarize class attributes
+        string = cleandoc(
+            f"""
+            Trace:
+            {indent}Attributes:
+            {indent*2}{len(self)} steps
+            {indent*2}{len(self.fluents)} fluents
+            {indent}Steps:
+            """
+        )
+        string += "\n"
+
+        # Dynamically get the spacing, 2n time
+        state_len = max([len(step.state.details()) for step in self]) + 4
+        string += f"{indent*2}{'Step':<5} {'State':^{state_len}} {'Action':<8}"
+        string += "\n"
+
+        # Create step string representation here, so formatting is consistent
+        for i, step in enumerate(self):
+            string += (
+                f"{indent*2}{i+1:<5} {step.state.details():<{state_len}} "
+                f"{step.action.details() if step.action is not None else '':<8}\n"
+            )
+
         return string
 
     def __len__(self):
@@ -70,7 +85,6 @@ def __getitem__(self, key: int):
 
     def __delitem__(self, key: int):
         del self.steps[key]
-        self.update()
 
     def __iter__(self):
         return iter(self.steps)
@@ -78,17 +92,17 @@ def __iter__(self):
     def __reversed__(self):
         return reversed(self.steps)
 
-    def __contains__(self, item: Step):
-        return item in self.steps
+    def __contains__(self, step: Step):
+        return step in self.steps
 
-    def append(self, item: Step):
-        self.steps.append(item)
-        self.__update_actions_and_fluents(item)
+    def append(self, step: Step):
+        self.steps.append(step)
+        self.__update_actions_and_fluents(step)
 
     def clear(self):
         self.steps.clear()
-        self.fluents = []
-        self.actions = []
+        self.fluents = set()
+        self.actions = set()
 
     def copy(self):
         return self.steps.copy()
@@ -123,156 +137,98 @@ def reverse(self):
     def sort(self, reverse: bool = False, key: Callable = lambda e: e.action.cost):
         self.steps.sort(reverse=reverse, key=key)
 
-    def __new_fluents_from_state(self, state: State):
-        """
-        Retrieves any new fluents (fluents not yet in this trace's list of fluents)
-        from the given state.
-
-        Args:
-            state (State): the state to extract new fluents from.
-        """
-        new = []
-        for fluent in state.fluents:
-            name = fluent.name
-            if name not in self.fluents and name not in new:
-                new.append(name)
-        return new
-
-    def __new_action_from_step(self, step: Step):
-        """
-        Retrieves the action from the given step if the action is new (not yet in this
-        trace's list of actions).
-
-        Args:
-            step (Step): the given step to extract the action from.
-        """
-        name = step.action.name
-        if name not in self.actions:
-            return name
-
     def __update_actions_and_fluents(self, step: Step):
-        """
-        Update this trace's actions and fluents after taking into account the action and
-        fluents of the step just added.
-
-        Args:
-            step (Step): the step just added to the trace.
-        """
-        new_fluents = self.__new_fluents_from_state(step.state)
-        if new_fluents:
-            self.fluents.extend(new_fluents)
-        new_act = self.__new_action_from_step(step)
-        if new_act:
-            self.actions.append(new_act)
+        self.fluents.update(step.state.keys())
+        new_act = step.action
+        if new_act is not None:
+            self.actions.add(new_act)
 
     def __reinit_actions_and_fluents(self):
-        self.fluents = []
-        self.actions = []
+        self.fluents = set()
+        self.actions = set()
         for step in self.steps:
             self.__update_actions_and_fluents(step)
 
-    def get_prev_states(self, action: Action):
-        """
-        Returns a list of the states of the trace before this action took place.
+    def get_pre_states(self, action: Action):
+        """Retrieves the list of states prior to the action in this trace.
 
-        Arguments
-        ---------
-        action : Action
-            An `Action` object.
+        Args:
+            action (Action):
+                The action to retrieve pre-states for.
 
-        Returns
-        -------
-        prev_states : List of States
-            A list of states before this action took place.
+        Returns:
+            The set of states prior to the action being performed in this
+            trace.
         """
-        prev_states = []
-        for step in self.steps:
+        prev_states = set()
+        for step in self:
             if step.action == action:
-                prev_states.append(step.state)
+                prev_states.add(step.state)
         return prev_states
 
     def get_post_states(self, action: Action):
-        """
-        Returns a list of the states of the trace after the given action took place.
+        """Retrieves the list of states after the action in this trace.
 
-        Arguments
-        ---------
-        action : Action
-            An `Action` object.
+        Args:
+            action (Action):
+                The action to retrieve post-states for.
 
-        Returns
-        -------
-        post_states : List of States
-            A list of states after this action took place.
+        Returns:
+            The set of states after the action was performed in this trace.
         """
-        post_states = []
-        for i in range(len(self.steps) - 1):
-            if self.steps[i].action == action:
-                post_states.append(self.steps[i + 1].state)
+        post_states = set()
+        for i, step in enumerate(self):
+            if step.action == action:
+                post_states.add(self[i + 1].state)
         return post_states
 
-    def get_sas_triples(self, action: Action):
-        """
-        Returns a list of tuples where each tuple contains the state of the trace
-        before the action, the action, and the state of the trace after the action.
-
-        Arguments
-        ---------
-        action : Action
-            An `Action` object.
-
-        Returns
-        -------
-        sas_triples : List of tuples
-            A list of tuples in the format (previous state, action, post-state).
-        """
-        sas_triples = []
-        triple = []
-        for i in range(len(self.steps)):
-            if self.steps[i].action == action:
-                triple.append(self.steps[i].state)
-                triple.append(action)
-                if i + 1 < len(self.steps):
-                    triple.append(self.steps[i + 1].state)
-                triple = tuple(triple)
-                sas_triples.append(triple)
-                triple = []
+    def get_sas_triples(self, action: Action) -> Set[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
+        the post-state.
+
+        Args:
+            action (Action):
+                The action to retrieve (S,A,S') triples for.
+
+        Returns:
+            A `SAS` object, containing the `pre_state`, `action`, and
+            `post_state`.
+        """
+        sas_triples = set()
+        for i, step in enumerate(self):
+            if step.action == action:
+                triple = SAS(step.state, action, self[i + 1].state)
+                sas_triples.add(triple)
         return sas_triples
 
     def get_total_cost(self):
-        """
-        Returns the total cost of all actions used in this Trace.
+        """Calculates the total cost of this trace.
 
-        Returns
-        -------
-        sum : int
-            The total cost of all actions used.
+        Returns:
+            The total cost of all actions performed in the trace.
         """
         sum = 0
         for step in self.steps:
             sum += step.action.cost
         return sum
 
-    def get_cost_range(self, start: int, end: int):
-        """
-        Returns the total cost of the actions in the specified range of this Trace.
-        The range starts at 1 and ends at the number of steps.
-
-        Arguments
-        ---------
-        start : int
-            The start of the specified range.
-        end: int
-            The end of the specified range.
-
-        Returns
-        -------
-        sum : int
-            The total cost of all actions in the specified range.
+    def get_slice_cost(self, start: int, end: int):
+        """Calculates the total cost of a slice of this trace.
+
+        Args:
+            start (int):
+                The start of the slice range. 1 <= start <= end.
+            end: int
+                The end of the slice range. start <= end <= n, where n is the
+                length of the trace (number of steps).
+
+        Returns:
+            The total cost of the slice of the trace.
         """
 
-        if start < 1 or end < 1 or start > len(self.steps) or end > len(self.steps):
+        if start < 1 or end < 1 or start > len(self) or end > len(self):
             raise self.InvalidCostRange(
                 "Range supplied goes out of the feasible range."
             )
@@ -286,38 +242,51 @@ def get_cost_range(self, start: int, end: int):
             sum += self.steps[i].action.cost
         return sum
 
-    def get_usage(self, action: Action):
-        """
-        Returns the percentage of the number of times this action was used compared to the total
-        number of actions taken.
-
-        Arguments
-        ---------
-        action : Action
-            An `Action` object.
-
-        Returns
-        -------
-        percentage : float
-            The percentage of the number of times this action was used compared to the total
-            number of actions taken.
+    def get_steps(self, action: Action):
+        """Retrieves all the Steps in the trace that use the specified action.
+
+        Args:
+            action (Action):
+                The Action to retrieve corresponding steps for.
+
+        Returns:
+            The set of steps that use the specified action.
+
         """
-        sum = 0
-        for step in self.steps:
+        steps = set()
+        for step in self:
             if step.action == action:
-                sum += 1
-        return sum / len(self.steps)
+                steps.add(action)
+        return steps
 
-    def tokenize(self, Token: Type[Observation], **kwargs):
-        """
-        Creates the observation tokens using the token provided by the Observation.
+    def get_usage(self, action: Action):
+        """Calculates how often an action was performed in this trace.
 
-        Arguments
-        ---------
-        Token : Observation subclass
-            An `Observation` subclass.
+        Args:
+            action (Action):
+                The action to find the usage of.
+
+        Returns:
+            The frequency of the action in this trace. A percentage, calculated
+            as the number of occurences of the action divided by the length of
+            the trace (number of steps).
         """
+        return len(self.get_steps(action)) / len(self)
+
+    def tokenize(self, Token: Type[Observation], **kwargs):
+        """Tokenizes the steps in this trace.
 
+        Args:
+            Token (Observation):
+                A subclass of `Observation`, defining the method of tokenization
+                for the steps.
+
+        Returns:
+            A list of observation tokens, corresponding to the steps in the
+            trace.
+        """
+        observations: List[Observation] = []
         for step in self.steps:
             token = Token(step=step, **kwargs)
-            self.observations.append(token)
+            observations.append(token)
+        return observations
diff --git a/macq/trace/trace_list.py b/macq/trace/trace_list.py
index e51cc829..913c5369 100644
--- a/macq/trace/trace_list.py
+++ b/macq/trace/trace_list.py
@@ -1,13 +1,25 @@
-from typing import Iterable, Union, List, Callable
-from . import Action
-from . import Trace
+from collections import defaultdict
+from typing import List, Callable, Type, Set
+from . import Action, Trace
+from ..observation import Observation
 
 
 class TraceList:
+    """A collection of traces.
+
+    A `list`-like object, where each element is a `Trace` of the same planning
+    problem.
+
+
+    Attributes:
+        traces (list):
+            The list of `Trace` objects.
+        generator (function | None):
+            The function used to generate the traces.
     """
-    A TraceList object is a list-like object that holds information about a
-    series of traces.
-    """
+
+    # Allow child classes to have traces as a list of any type
+    traces: List
 
     class MissingGenerator(Exception):
         def __init__(
@@ -21,44 +33,20 @@ def __init__(
 
     def __init__(
         self,
-        traces: Union[List[Trace], None] = None,
-        generator: Union[Callable, None] = None,
+        traces: List[Trace] = None,
+        generator: Callable = None,
     ):
-        """
-        Creates a TraceList object. This stores a list of traces and optionally
-        the function used to generate the traces.
+        """Initializes a TraceList with a list of traces and a generator.
 
-        Attributes
-        ----------
-        traces : list
-            The list of `Trace` objects.
-        generator : funtion | None
-            (Optional) The function used to generate the traces.
+        Args:
+            traces (list):
+                Optional; The list of `Trace` objects.
+            generator (function):
+                Optional; The function used to generate the traces.
         """
-        self.traces: List[Trace] = [] if traces is None else traces
+        self.traces = traces if traces is not None else []
         self.generator = generator
 
-
-    def __str__(self):
-        string = "["
-        for trace in self:
-            string += "\n"
-            string += "-" * 100
-            string += "\n\n"
-            string += str(trace)
-        string += "]"
-        return string
-
-    def __repr__(self):
-        rep = "["
-        for trace in self:
-            rep += "\n"
-            rep += "-" * 100
-            rep += "\n\n"
-            rep += repr(trace)
-        rep += "]"
-        return rep
-
     def __len__(self):
         return len(self.traces)
 
@@ -77,10 +65,10 @@ def __iter__(self):
     def __reversed__(self):
         return reversed(self.traces)
 
-    def __contains__(self, item: Trace):
+    def __contains__(self, item):
         return item in self.traces
 
-    def append(self, item: Trace):
+    def append(self, item):
         self.traces.append(item)
 
     def clear(self):
@@ -89,22 +77,22 @@ def clear(self):
     def copy(self):
         return self.traces.copy()
 
-    def count(self, value: Trace):
+    def count(self, value):
         return self.traces.count(value)
 
-    def extend(self, iterable: Iterable[Trace]):
+    def extend(self, iterable):
         self.traces.extend(iterable)
 
-    def index(self, value: Trace):
+    def index(self, value):
         return self.traces.index(value)
 
-    def insert(self, index: int, item: Trace):
+    def insert(self, index: int, item):
         self.traces.insert(index, item)
 
     def pop(self):
         return self.traces.pop()
 
-    def remove(self, value: Trace):
+    def remove(self, value):
         self.traces.remove(value)
 
     def reverse(self):
@@ -113,14 +101,103 @@ def reverse(self):
     def sort(self, reverse: bool = False, key: Callable = lambda e: e.get_cost()):
         self.traces.sort(reverse=reverse, key=key)
 
+    def print(self):
+        string = "TraceList:\n"
+        for trace in self:
+            for line in trace.details().splitlines():
+                string += f"    {line}\n"
+        print(string)
+
     def generate_more(self, num: int):
+        """Generates more traces using the generator function.
+
+        Args:
+            num (int):
+                The number of additional traces to generate.
+
+        Raises:
+            MissingGenerator: Cannot generate more traces if the generator
+            function is not provided.
+        """
         if self.generator is None:
             raise self.MissingGenerator(self)
 
         self.traces.extend(self.generator(num))
 
     def get_usage(self, action: Action):
+        """Calculates how often an action was performed in each of the traces.
+
+        Args:
+            action (Action):
+                The action to find the usage of.
+
+        Returns:
+            The frequency of the action in each of the traces. A percentage,
+            calculated as the number of occurences of the action divided by the
+            length of the trace (number of steps).
+        """
         usages = []
         for trace in self:
             usages.append(trace.get_usage(action))
         return usages
+
+    def tokenize(self, Token: Type[Observation], **kwargs):
+        """Tokenizes the steps in this trace.
+
+        Args:
+            Token (Observation):
+                A subclass of `Observation`, defining the method of tokenization
+                for the steps.
+        """
+
+        return ObservationList(self, Token, **kwargs)
+
+
+class ObservationList(TraceList):
+    traces: List[List[Observation]]
+    # Disable methods
+    generate_more = property()
+    get_usage = property()
+    tokenize = property()
+
+    def __init__(self, traces: TraceList, Token: Type[Observation], **kwargs):
+        super(ObservationList, self).__init__()
+        self.type = Token
+        for i, trace in enumerate(traces):
+            tokens = trace.tokenize(Token, trace_num=i, **kwargs)
+            self.append(tokens)
+
+    def fetch_observations(self, query: dict):
+        matches = list()
+        trace: List[Observation]
+        for i, trace in enumerate(self):
+            matches.append(set())
+            for obs in trace:
+                if obs.matches(query):  # if no matches, set can be empty
+                    matches[i].add(obs)
+        return matches  # list of sets of matching fluents from each trace
+
+    def fetch_observation_windows(self, query: dict, left: int, right: int):
+        windows = []
+        matches = self.fetch_observations(query)
+        trace: Set[Observation]
+        for i, trace in enumerate(matches):  # i corresponds to trace index in self
+            for obs in trace:
+                start = obs.index - left
+                end = obs.index + right + 1
+                windows.append(self[i][start:end])
+        return windows
+
+    def get_transitions(self, action: Action):
+        query = {"action": action.details()}
+        return self.fetch_observation_windows(query, 0, 1)
+
+    def get_all_transitions(self):
+        actions = set()
+        for trace in self:
+            for obs in trace:
+                action = obs.action
+                if action is not None:
+                    actions.add(action)
+
+        return dict(map(lambda action: (action, self.get_transitions(action)), actions))
diff --git a/macq/utils/__init__.py b/macq/utils/__init__.py
index bd054a23..6506c785 100644
--- a/macq/utils/__init__.py
+++ b/macq/utils/__init__.py
@@ -1,3 +1,3 @@
 from .timer import set_timer, TraceSearchTimeOut
 
-__all__ = ["set_timer", "TraceSearchTimeOut"]
\ No newline at end of file
+__all__ = ["set_timer", "TraceSearchTimeOut"]


Test patch:
diff --git a/tests/extract/extract.py b/tests/extract/extract.py
new file mode 100644
index 00000000..b07f37de
--- /dev/null
+++ b/tests/extract/extract.py
@@ -0,0 +1,9 @@
+import pytest
+from tests.test_core import generate_test_trace_list
+from macq.extract import Extract, modes
+
+
+def generate_observer_model():
+    traces = generate_test_trace_list(2)
+    model = Extract(traces, modes.OBSERVER)
+    return model
diff --git a/tests/extract/model.py b/tests/extract/model.py
new file mode 100644
index 00000000..d4f95984
--- /dev/null
+++ b/tests/extract/model.py
@@ -0,0 +1,10 @@
+import pytest
+from macq.model import Model
+from tests.extract.extract import generate_observer_model
+
+
+def test_model_serialization():
+    model = generate_observer_model()
+    s = model.serialize()
+    model2 = Model.deserialize(s)
+    assert model2.serialize() == s
diff --git a/tests/extract/test_observer.py b/tests/extract/test_observer.py
new file mode 100644
index 00000000..c41b6df9
--- /dev/null
+++ b/tests/extract/test_observer.py
@@ -0,0 +1,12 @@
+from macq.extract import Extract, modes
+from macq.observation import IdentityObservation
+from macq.trace import *
+from tests.utils.realistic_trace import real_trace_list
+
+
+if __name__ == "__main__":
+    traces = real_trace_list()
+    observations = traces.tokenize(IdentityObservation)
+    traces.print()
+    model = Extract(observations, modes.OBSERVER)
+    print(model.details())
diff --git a/tests/test_core.py b/tests/test_core.py
index 2476d2f5..5cc12514 100644
--- a/tests/test_core.py
+++ b/tests/test_core.py
@@ -1,18 +1,24 @@
-from macq.trace import CustomObject, Fluent, Action, Step, State, Trace, TraceList
+import pytest
+from typing import List, Dict
+from macq.trace import (
+    PlanningObject,
+    Fluent,
+    Action,
+    Step,
+    State,
+    Trace,
+    SAS,
+    TraceList,
+)
 from macq.observation import IdentityObservation
 from pathlib import Path
 from macq.utils.timer import TraceSearchTimeOut
 from macq.generate.pddl import VanillaSampling
-from macq.generate.pddl.planning_domains_api import get_problem
 
 InvalidCostRange = Trace.InvalidCostRange
-InvalidFluent = Action.InvalidFluent
+# InvalidFluent = Action.InvalidFluent
 MissingGenerator = TraceList.MissingGenerator
 
-from typing import List
-from pathlib import Path
-import pytest
-
 # HELPER FUNCTIONS
 
 
@@ -31,19 +37,34 @@ def generate_test_fluents(num_fluents: int):
         The list of testing fluents generated.
     """
     fluents = []
-    objects = [CustomObject("number", str(o)) for o in range(num_fluents)]
+    objects = [PlanningObject("number", str(o)) for o in range(num_fluents)]
     for i in range(num_fluents):
         fluent_name = "fluent" + " " + str(i + 1)
-        if i % 2 == 0:
-            value = True
-        else:
-            value = False
-        fluent = Fluent(fluent_name, [objects[i]], value)
+        fluent = Fluent(fluent_name, [objects[i]])
         fluents.append(fluent)
     return fluents
 
 
-def generate_test_actions(num_actions: int):
+def generate_fluent_dicts(fluents: List[Fluent]):
+    """
+    Generates fluent dictionaries to be used for testing.
+
+    Args:
+        fluents (List[Fluent]): The list of fluents to create dictionaries from.
+
+    Returns:
+        fluents_dict (dict): The dictionary of fluent names mapped to boolean values.
+    """
+    fluents_dict = {}
+    for i in range(len(fluents)):
+        if i % 2 == 0:
+            fluents_dict[fluents[i]] = True
+        else:
+            fluents_dict[fluents[i]] = False
+    return fluents_dict
+
+
+def generate_test_actions(num_actions: int, objects: List[PlanningObject]):
     """
     Generates basic actions to be used for testing.
 
@@ -51,18 +72,19 @@ def generate_test_actions(num_actions: int):
     ---------
     num_actions : int
         The number of actions to generate.
+    objects: List of CustomObjects
+        The objects available to these actions.
 
     Returns
     -------
     actions : List of Actions
         The list of testing actions generated.
     """
-    objects = [CustomObject("number", str(o)) for o in range(num_actions)]
     actions = []
     for i in range(num_actions):
-        action_name = "action " + str(i + 1)
+        action_name = "action" + " " + str(i + 1)
         # action 1 has a cost of 1, etc.
-        action = Action(action_name, objects, [], [], [], i + 1)
+        action = Action(action_name, objects, i + 1)
         actions.append(action)
     return actions
 
@@ -88,7 +110,7 @@ def get_fluent_obj(fluents: List[Fluent]):
     return objects
 
 
-def generate_test_states(num_states: int):
+def generate_test_states(num_states: int, fluents: Dict[Fluent, bool] = {}):
     """
     Generate states to be used for testing, using the given fluents (each state will add a fluent)
 
@@ -96,6 +118,8 @@ def generate_test_states(num_states: int):
     ---------
     num_states : int
         The number of states to generate.
+    fluents : dict of Fluents
+        The fluents that will be used to make up the states.
 
     Returns
     -------
@@ -103,17 +127,17 @@ def generate_test_states(num_states: int):
         The list of testing states generated.
     """
     states = []
-    next_fluents = []
-    fluents = generate_test_fluents(num_states)
+    next_fluents = {}
     for i in range(num_states):
-        state_name = "state " + str(i + 1)
-        next_fluents = fluents[: i + 1]
-        state = State(next_fluents)
+        if i < len(list(fluents.keys())):
+            current_fluent = list(fluents.keys())[i]
+            next_fluents[current_fluent] = fluents[current_fluent]
+        state = State(next_fluents.copy())
         states.append(state)
     return states
 
 
-def generate_test_steps(num_steps: int):
+def generate_test_steps(num_steps: int, actions: List[Action], states: List[State]):
     """
     Generate steps to be used for testing, given the number of steps and possible actions and states.
 
@@ -121,6 +145,10 @@ def generate_test_steps(num_steps: int):
     ---------
     num_steps : int
         The number of steps to generate.
+    actions : List of Actions
+        The list of possible actions to be used for the generated steps.
+    states : List of States
+        The list of possible states to be used for the generated steps.
 
     Returns
     -------
@@ -128,10 +156,20 @@ def generate_test_steps(num_steps: int):
         The list of testing steps generated.
     """
     steps = []
-    actions = generate_test_actions(num_steps)
-    states = generate_test_states(num_steps)
+    # indices for actions and states respectively
+    a_index = 0
+    s_index = 0
     for i in range(num_steps):
-        step = Step(actions[i], states[i], i + 1)
+        step = Step(states[a_index], actions[s_index], i)
+        # cycle through actions and states
+        if a_index < len(actions):
+            a_index += 1
+        else:
+            a_index = 0
+        if s_index < len(states):
+            s_index += 1
+        else:
+            s_index = 0
         steps.append(step)
     return steps
 
@@ -150,114 +188,108 @@ def generate_test_trace(complexity: int):
     trace : Trace
         The testing trace generated.
     """
-    trace = Trace(generate_test_steps(complexity))
+    fluents = generate_test_fluents(complexity)
+    actions = generate_test_actions(complexity, get_fluent_obj(fluents))
+    fluents_dict = generate_fluent_dicts(fluents)
+    states = generate_test_states(complexity, fluents_dict)
+    steps = generate_test_steps(complexity, actions, states)
+    trace = Trace(steps)
     return trace
 
 
 # TESTS FOR ACTION CLASS
 
-# ensure that invalid fluents can't be added to actions
+# ensure that invalid fluents can't be added to actions.
+# NOTE: DEBATING RAISING AN ERROR VS. JUST PRINTING A WARNING.
+"""
 def test_action_errors():
-    objects = [CustomObject("number", str(o)) for o in range(6)]
+    objects = [PlanningObject("number", str(o)) for o in range(6)]
     action = Action("put down", objects, [], [], [], 1)
-    other = CustomObject("other", 10)
-    fluent_other = Fluent("put down other", [other], True)
+    other = PlanningObject("other", "10")
+    fluent_other = Fluent("put down other", [other])
 
     with pytest.raises(InvalidFluent):
-        action.add_precond([fluent_other])
-        action.add_effect_add([fluent_other])
-        action.add_effect_delete([fluent_other])
-
-
-# ensure that valid fluents can be added as action preconditions
-def test_action_add_preconditions():
-    fluents = generate_test_fluents(3)
-    (fl1, fl2, fl3) = tuple(fluents)
-    action = Action("put down", get_fluent_obj(fluents), [], [], [], 1)
+        action.update_precond([fluent_other])
+        action.update_add([fluent_other])
+        action.update_delete([fluent_other])
+"""
 
-    action.add_precond([fl1])
-    assert action.precond == [fl1]
-    action.add_precond([fl2, fl3])
-    assert action.precond == [fl1, fl2, fl3]
 
+# ensure that valid object parameters can be added and subsequently referenced
+# def test_action_add_params():
+#     objects = [PlanningObject("number", str(o)) for o in range(6)]
+#     action = Action("put down", objects, 1)
+#     other = PlanningObject("other", "other")
+#     fluent_other = Fluent("put down other", [other])
 
-# ensure that valid fluents can be added as action effects
-def test_action_add_effects():
-    fluents = generate_test_fluents(6)
-    (fl1, fl2, fl3, fl4, fl5, fl6) = tuple(fluents)
-    action = Action("put down", get_fluent_obj(fluents), [], [], [], 1)
+#     action.add_parameter(other)
+# action.update_precond([fluent_other])
+# action.update_add([fluent_other])
+# action.update_delete([fluent_other])
+# assert action.precond == {fluent_other}
+# assert action.add == {fluent_other}
+# assert action.delete == {fluent_other}
 
-    action.add_effect_add([fl4])
-    assert action.add == [fl4]
-    action.add_effect_add([fl5, fl6])
-    assert action.add == [fl4, fl5, fl6]
-    action.add_effect_delete([fl1])
-    assert action.delete == [fl1]
-    action.add_effect_delete([fl2, fl3])
-    assert action.delete == [fl1, fl2, fl3]
 
+# TESTS FOR TRACE CLASS
 
-# ensure that valid object parameters can be added and subsequently referenced
-def test_action_add_params():
+# test the functionality to add steps to a trace (NOT YET WORKING -- need equality dunders). pass for now
+def test_trace_add_steps():
+    """
     objects = [CustomObject("number", str(o)) for o in range(6)]
     action = Action("put down", objects, [], [], [], 1)
-    other = CustomObject("other", 10)
-    fluent_other = Fluent("put down other", [other], True)
-
-    action.add_parameter(other)
-    action.add_precond([fluent_other])
-    action.add_effect_add([fluent_other])
-    action.add_effect_delete([fluent_other])
-    assert action.precond == [fluent_other]
-    assert action.add == [fluent_other]
-    assert action.delete == [fluent_other]
+    fluent = Fluent("on table", [objects[0]], True)
+    state = State([fluent])
+    trace = generate_test_trace(3)
+    step4 = generate_test_steps(1, [action], [state])
 
+    result = copy.deepcopy(trace.steps)
+    result.append(step4)
 
-# TESTS FOR TRACE CLASS
-
-# ensure that the Trace base_fluents() and base_actions() functions work correctly
-def test_trace_base():
-    trace = generate_test_trace(3)
-    print(trace.fluents)
-    print(trace.actions)
-    assert trace.fluents == ["fluent 1", "fluent 2", "fluent 3"]
-    assert trace.actions == ["action 1", "action 2", "action 3"]
+    trace.add_steps([step4])
+    assert trace.steps == result
+    """
+    pass
 
 
 # test that the previous states are being retrieved correctly
-def test_trace_prev_states():
+def test_trace_pre_states():
     trace = generate_test_trace(3)
     # get the first and last action
-    (action1, action3) = (trace[0].action, trace[2].action)
+    (action1, action2) = (trace[0].action, trace[1].action)
     # get the first and last state
-    (state1, state3) = (trace[0].state, trace[2].state)
+    (state1, state2) = (trace[0].state, trace[1].state)
 
-    assert trace.get_prev_states(action1) == [state1]
-    assert trace.get_prev_states(action3) == [state3]
+    assert isinstance(action1, Action)
+    assert isinstance(action2, Action)
+    assert trace.get_pre_states(action1) == {state1}
+    assert trace.get_pre_states(action2) == {state2}
 
 
 # test that the post states are being retrieved correctly
 def test_trace_post_states():
     trace = generate_test_trace(3)
-    # get the first and last action
-    (action1, action3) = (trace[0].action, trace[2].action)
-    # get the second state
-    state2 = trace[1].state
+    # get the first and second action
+    (action1, action2) = (trace[0].action, trace.steps[1].action)
+    # get the second and last state
+    (state2, state3) = (trace[1].state, trace[2].state)
 
-    assert trace.get_post_states(action1) == [state2]
-    assert trace.get_post_states(action3) == []
+    assert isinstance(action1, Action)
+    assert isinstance(action2, Action)
+    assert trace.get_post_states(action1) == {state2}
+    assert trace.get_post_states(action2) == {state3}
 
 
 # test trace SAS triples function
 def test_trace_get_sas_triples():
     trace = generate_test_trace(3)
-    # get the second and last action
-    (action2, action3) = (trace[1].action, trace[2].action)
+    # get the second action
+    action2 = trace.steps[1].action
     # get the second and last state
-    (state2, state3) = (trace[1].state, trace[2].state)
+    (state2, state3) = (trace.steps[1].state, trace.steps[2].state)
 
-    assert trace.get_sas_triples(action2) == [(state2, action2, state3)]
-    assert trace.get_sas_triples(action3) == [(state3, action3)]
+    assert isinstance(action2, Action)
+    assert trace.get_sas_triples(action2) == {SAS(state2, action2, state3)}
 
 
 # test that the total cost is working correctly
@@ -269,72 +301,62 @@ def test_trace_total_cost():
 # test that the cost range is working correctly
 def test_trace_valid_cost_range():
     trace = generate_test_trace(5)
-    assert trace.get_cost_range(1, 3) == 6
-    assert trace.get_cost_range(2, 3) == 5
-    assert trace.get_cost_range(1, 5) == 15
-    assert trace.get_cost_range(4, 5) == 9
+    assert trace.get_slice_cost(1, 3) == 6
+    assert trace.get_slice_cost(2, 3) == 5
+    assert trace.get_slice_cost(1, 5) == 15
+    assert trace.get_slice_cost(4, 5) == 9
 
 
 # test that incorrect provided cost ranges throw errors
 def test_trace_invalid_cost_range():
     trace = generate_test_trace(3)
     with pytest.raises(InvalidCostRange):
-        trace.get_cost_range(3, 1)
-        trace.get_cost_range(0, 2)
-        trace.get_cost_range(1, 5)
+        trace.get_slice_cost(3, 1)
+        trace.get_slice_cost(0, 2)
+        trace.get_slice_cost(1, 5)
 
 
 # test trace action usage
 def test_trace_usage():
     trace = generate_test_trace(3)
     # get the first action
-    action1 = trace[0].action
+    action1 = trace.steps[0].action
+    assert isinstance(action1, Action)
     assert trace.get_usage(action1) == 1 / 3
 
 
 # test trace tokenize function
 def test_trace_tokenize():
     trace = generate_test_trace(3)
-    (step1, step2, step3) = (trace[0], trace[1], trace[2])
-    trace.tokenize(IdentityObservation)
-    print(trace.observations)
-    print(
-        [
-            IdentityObservation(step1),
-            IdentityObservation(step2),
-            IdentityObservation(step3),
-        ]
-    )
-    assert trace.observations == [
-        IdentityObservation(step1),
-        IdentityObservation(step2),
-        IdentityObservation(step3),
-    ]
+    observations = trace.tokenize(IdentityObservation)
+    for i, obs in enumerate(observations):
+        assert obs == IdentityObservation(step=trace[i])
     # test equality dunder by attempting to compare an object of a different type
-    assert trace.observations != step1
+    assert observations != trace[0]
 
 
 # test the timer wrapper on vanilla trace generation
+# def test_timer_vanilla_wrapper():
+#     # exit out to the base macq folder so we can get to /tests
+#     base = Path(__file__).parent.parent
+#     dom = (base / "tests/pddl_testing_files/playlist_domain.pddl").resolve()
+#     prob = (base / "tests/pddl_testing_files/playlist_problem.pddl").resolve()
 
+#     with pytest.raises(TraceSearchTimeOut):
+#         VanillaSampling(dom, prob, 10, 5)
 
-def test_timer_wrapper_vanilla():
-    # exit out to the base macq folder so we can get to /tests
-    base = Path(__file__).parent.parent
-    dom = (base / "tests/pddl_testing_files/playlist_domain.pddl").resolve()
-    prob = (base / "tests/pddl_testing_files/playlist_problem.pddl").resolve()
 
-    with pytest.raises(TraceSearchTimeOut):
-        vanilla = VanillaSampling(dom=dom, prob=prob, plan_len=10, num_traces=5)
-
-
-# generate testing trace lists
 def generate_test_trace_list(length: int):
-    trace = generate_test_trace(3)
-    traces = [trace] * length
+    from random import randint
+
+    traces = []
+    for _ in range(length):
+        comp = randint(1, 3)
+        trace = generate_test_trace(comp)
+        traces.append(trace)
     return TraceList(traces)
 
 
-# test trace lists
 def test_trace_list():
     trace_list = generate_test_trace_list(5)
 
@@ -353,89 +375,3 @@ def test_trace_list():
     usages = trace_list.get_usage(action)
     for i, trace in enumerate(trace_list):
         assert usages[i] == trace.get_usage(action)
-
-
-# test trace append function
-def test_trace_append():
-    trace = generate_test_trace(3)
-    steps = generate_test_steps(4)
-    trace.append(steps[3])
-    assert trace.fluents == ["fluent 1", "fluent 2", "fluent 3", "fluent 4"]
-    assert trace.actions == ["action 1", "action 2", "action 3", "action 4"]
-    # assert trace.steps == steps
-
-
-# test trace clear function
-def test_trace_clear():
-    trace = generate_test_trace(3)
-    trace.clear()
-    assert trace.fluents == []
-    assert trace.actions == []
-    assert trace.steps == []
-
-
-# test trace extend function
-def test_trace_extend():
-    trace = generate_test_trace(3)
-    steps = generate_test_steps(7)
-    trace.extend(steps[3:])
-    assert trace.fluents == [
-        "fluent 1",
-        "fluent 2",
-        "fluent 3",
-        "fluent 4",
-        "fluent 5",
-        "fluent 6",
-        "fluent 7",
-    ]
-    assert trace.actions == [
-        "action 1",
-        "action 2",
-        "action 3",
-        "action 4",
-        "action 5",
-        "action 6",
-        "action 7",
-    ]
-    # assert trace.steps == steps
-
-
-# test trace insert function
-def test_trace_insert():
-    trace = generate_test_trace(3)
-    steps = generate_test_steps(4)
-    trace.insert(0, steps[3])
-    assert trace.fluents == ["fluent 1", "fluent 2", "fluent 3", "fluent 4"]
-    assert trace.actions == ["action 1", "action 2", "action 3", "action 4"]
-    # assert trace.steps == [steps[3], steps[0], steps[1], steps[2]]
-
-
-# test trace pop function
-def test_trace_pop():
-    trace = generate_test_trace(3)
-    steps = trace.steps.copy()
-    trace.pop()
-    assert trace.fluents == ["fluent 1", "fluent 2"]
-    assert trace.actions == ["action 1", "action 2"]
-    # assert trace == steps[:-1]
-
-
-# test trace remove function
-def test_trace_remove():
-    trace = generate_test_trace(3)
-    steps = trace.steps.copy()
-    trace.remove(steps[1])
-    assert trace.fluents == ["fluent 1", "fluent 2", "fluent 3"]
-    assert trace.actions == ["action 1", "action 3"]
-    # assert trace.steps == [steps[0], steps[2]]
-
-
-if __name__ == "__main__":
-    # exit out to the base macq folder so we can get to /tests
-    base = Path(__file__).parent.parent
-    dom = (base / "tests/pddl_testing_files/playlist_domain.pddl").resolve()
-    prob = (base / "tests/pddl_testing_files/playlist_problem.pddl").resolve()
-    vanilla = VanillaSampling(dom=dom, prob=prob, plan_len=3, num_traces=2)
-    # print(vanilla.traces)
-    # vanilla = VanillaSampling(problem_id=123, plan_len=3, num_traces=1)
-    print(vanilla.traces)
diff --git a/tests/utils/realistic_trace.py b/tests/utils/realistic_trace.py
new file mode 100644
index 00000000..6b15105b
--- /dev/null
+++ b/tests/utils/realistic_trace.py
@@ -0,0 +1,185 @@
+from macq.trace import *
+
+
+def real_trace_list():
+    # Objects
+    blockA = PlanningObject("block", "A")
+    blockB = PlanningObject("block", "B")
+
+    # Fluents
+    a_clear = Fluent("clear", [blockA])
+    b_clear = Fluent("clear", [blockB])
+    a_on_table = Fluent("on table", [blockA])
+    b_on_table = Fluent("on table", [blockB])
+    holding_a = Fluent("holding", [blockA])
+    holding_b = Fluent("holding", [blockB])
+    a_on_b = Fluent("on", [blockA, blockB])
+    b_on_a = Fluent("on", [blockA, blockB])
+
+    # Actions
+    pick_up_a = Action("pick up", [blockA])
+    pick_up_b = Action("pick up", [blockB])
+    stack_a_b = Action("stack", [blockA, blockB])
+    stack_b_a = Action("stack", [blockB, blockA])
+
+    # Traces
+    t1 = Trace(
+        [
+            Step(
+                State(
+                    {
+                        a_clear: True,
+                        b_clear: True,
+                        a_on_table: True,
+                        b_on_table: True,
+                        holding_a: False,
+                        holding_b: False,
+                        a_on_b: False,
+                        b_on_a: False,
+                    }
+                ),
+                pick_up_a,
+                0,
+            ),
+            Step(
+                State(
+                    {
+                        a_clear: True,
+                        b_clear: True,
+                        a_on_table: False,
+                        b_on_table: True,
+                        holding_a: True,
+                        holding_b: False,
+                        a_on_b: False,
+                        b_on_a: False,
+                    }
+                ),
+                stack_a_b,
+                1,
+            ),
+            Step(
+                State(
+                    {
+                        a_clear: True,
+                        b_clear: False,
+                        a_on_table: False,
+                        b_on_table: True,
+                        holding_a: False,
+                        holding_b: False,
+                        a_on_b: True,
+                        b_on_a: False,
+                    }
+                ),
+                None,
+                2,
+            ),
+        ]
+    )
+
+    t2 = Trace(
+        [
+            Step(
+                State(
+                    {
+                        a_clear: True,
+                        b_clear: True,
+                        a_on_table: True,
+                        b_on_table: True,
+                        holding_a: False,
+                        holding_b: False,
+                        a_on_b: False,
+                        b_on_a: False,
+                    }
+                ),
+                pick_up_b,
+                0,
+            ),
+            Step(
+                State(
+                    {
+                        a_clear: True,
+                        b_clear: True,
+                        a_on_table: True,
+                        b_on_table: False,
+                        holding_a: False,
+                        holding_b: True,
+                        a_on_b: False,
+                        b_on_a: False,
+                    }
+                ),
+                stack_b_a,
+                1,
+            ),
+            Step(
+                State(
+                    {
+                        a_clear: False,
+                        b_clear: True,
+                        a_on_table: True,
+                        b_on_table: False,
+                        holding_a: False,
+                        holding_b: False,
+                        a_on_b: False,
+                        b_on_a: True,
+                    }
+                ),
+                None,
+                2,
+            ),
+        ]
+    )
+
+    t3 = Trace(
+        [
+            Step(
+                State(
+                    {
+                        a_clear: True,
+                        b_clear: False,
+                        a_on_table: False,
+                        b_on_table: True,
+                        holding_a: False,
+                        holding_b: False,
+                        a_on_b: True,
+                        b_on_a: False,
+                    }
+                ),
+                pick_up_a,
+                0,
+            ),
+            Step(
+                State(
+                    {
+                        a_clear: True,
+                        b_clear: True,
+                        a_on_table: False,
+                        b_on_table: True,
+                        holding_a: True,
+                        holding_b: False,
+                        a_on_b: False,
+                        b_on_a: False,
+                    }
+                ),
+                stack_a_b,
+                1,
+            ),
+            Step(
+                State(
+                    {
+                        a_clear: True,
+                        b_clear: False,
+                        a_on_table: False,
+                        b_on_table: True,
+                        holding_a: False,
+                        holding_b: False,
+                        a_on_b: True,
+                        b_on_a: False,
+                    }
+                ),
+                None,
+                2,
+            ),
+        ]
+    )
+
+    return TraceList([t1, t2, t1, t3])