HorizonEncoder

HorizonEncoder

Bases: HGraphEncoder

Horizon lookahead encoder backed by HorizonHGraphEncoderEngine.

This encoder combines a root state, a TransitionDAG and goals into one hetero graph representation.

Source code in src/mifrost/encoders/horizon.py

class HorizonEncoder(HGraphEncoder):
    """
    Horizon lookahead encoder backed by ``HorizonHGraphEncoderEngine``.

    This encoder combines a root state, a ``TransitionDAG`` and goals into one
    hetero graph representation.
    """

    def __init__(
        self,
        domain: DomainInput,
        *,
        backend: HorizonBackendName | str | None = None,
        transition_mode: HorizonEncoderMode | str | None = None,
        target_symbol_prefix: str | None = None,
        target_sources: Iterable[TargetSource | str] | None = None,
        parent_relation: str | None = None,
        sibling_relation: str | None = None,
        cousin_relation: str | None = None,
        enable_parent_relation: bool | None = None,
        enable_sibling_relation: bool | None = None,
        enable_cousin_relation: bool | None = None,
        root_policy: RootPolicy | str | None = None,
        max_goal_level: int | None = None,
        symbol_type_id: str | None = DEFAULT_SYMBOL_TYPE_ID,
        ignore_actions: bool | None = None,
        add_nullary_predicates: bool | None = None,
        include_lgan_edges: bool | None = None,
        include_static: bool | None = None,
        include_empty_edge_types: bool | None = None,
        export_node_names: bool | None = None,
        support_literals: bool | None = None,
        goal_derivations: Iterable[Any] | None = None,
        nullary_object_name: str | None = None,
        lgan_tn_edge_pos: str | None = DEFAULT_LGAN_TN_EDGE_POS,
        lgan_nn_edge_pos: str | None = DEFAULT_LGAN_NN_EDGE_POS,
        lgan_rr_edge_pos: str | None = DEFAULT_LGAN_RR_EDGE_POS,
        history_link_relation: str | None = DEFAULT_HISTORY_LINK_RELATION,
    ) -> None:
        """Create a hetero horizon encoder.

        This lane reads a root state plus a `TransitionDAG` and creates
        candidate state rows.

        `target_sources` is therefore simple on this lane:

        - `state`: successor or candidate states from the DAG

        `root_policy` controls how the root state is treated:

        - `include`: encode the root and expose it as a target
        - `encode_only`: encode the root, but omit it from targets and metadata
        - `exclude`: omit the root from targets and keep root facts as base facts

        The main-lane sources `action`, `goal`, `subgoal`, and `history` do
        not create separate targets here. When `include_lgan_edges=True`, LGAN
        anchors are those candidate state rows. There is no
        `lgan_anchor_sources` switch here.
        """
        normalized_root_policy = normalize_root_policy(root_policy)
        config = self._make_config(
            _HORIZON_CONFIG_CLS,
            symbol_type_id=symbol_type_id,
            target_symbol_prefix=target_symbol_prefix,
            ignore_actions=ignore_actions,
            add_nullary_predicates=add_nullary_predicates,
            include_lgan_edges=include_lgan_edges,
            include_static=include_static,
            include_empty_edge_types=include_empty_edge_types,
            export_node_names=export_node_names,
            max_goal_level=max_goal_level,
            support_literals=support_literals,
            goal_derivations=goal_derivations,
            nullary_object_name=nullary_object_name,
            lgan_tn_edge_pos=lgan_tn_edge_pos,
            lgan_nn_edge_pos=lgan_nn_edge_pos,
            lgan_rr_edge_pos=lgan_rr_edge_pos,
            history_link_relation=history_link_relation,
            transition_mode=_normalize_horizon_mode(transition_mode),
            target_sources=(
                None
                if target_sources is None
                else self._normalize_horizon_target_sources(target_sources)
            ),
            parent_relation=parent_relation,
            sibling_relation=sibling_relation,
            cousin_relation=cousin_relation,
            enable_parent_relation=enable_parent_relation,
            enable_sibling_relation=enable_sibling_relation,
            enable_cousin_relation=enable_cousin_relation,
            root_policy=normalized_root_policy,
        )
        self._runtime = create_horizon_runtime(domain, config, backend=backend)
        self._engine = self._runtime.engine
        self._config = self._engine.config
        self.backend = self._runtime.backend_name
        self.symbol_type_id = self._config.symbol_type_id
        self.target_symbol_prefix = self._config.target_symbol_prefix
        self.parent_relation = self._config.parent_relation
        self.sibling_relation = self._config.sibling_relation
        self.cousin_relation = self._config.cousin_relation
        self.lgan_tn_edge_pos = self._config.lgan_tn_edge_pos
        self.lgan_nn_edge_pos = self._config.lgan_nn_edge_pos
        self.lgan_rr_edge_pos = self._config.lgan_rr_edge_pos
        self.include_lgan_edges = self._config.include_lgan_edges
        self.lgan_anchor_sources = set(self._config.lgan_anchor_sources)
        self._lgan_edge_positions = {
            self.lgan_tn_edge_pos,
            self.lgan_nn_edge_pos,
            self.lgan_rr_edge_pos,
        }

    @staticmethod
    def _normalize_horizon_target_sources(
        target_sources: Iterable[TargetSource | str],
    ) -> set[TargetSource]:
        from ._target_sources import normalize_target_sources

        normalized = normalize_target_sources(target_sources)
        assert normalized is not None
        return normalized

    @property
    def engine(self) -> Any:
        """Expose the underlying C++ horizon engine."""
        return self._engine

    def _encode(
        self,
        root: StateInput,
        dag: TransitionDAG | RXStateDAG | None = None,
        *,
        goals: GoalBatchInput = None,
        actions: ActionBatchInput = None,
        subgoal_layers: SubgoalLayersInput = None,
        history_subgoals: HistorySubgoalInput | None = None,
        history_max_steps: int | None = None,
        **_,
    ) -> BatchEncoding:
        """Encode one root/DAG pair."""
        validate_single_unsupported_lanes(
            actions=actions,
            history_subgoals=history_subgoals,
            history_max_steps=history_max_steps,
        )
        return self._runtime.encode(
            root, dag, goals=goals, subgoal_layers=subgoal_layers
        )

    def _encode_one_into_builder(
        self,
        root: StateInput,
        builder: Any,
        *,
        dag: TransitionDAG | RXStateDAG | None = None,
        goals: GoalBatchInput = None,
        actions: ActionBatchInput = None,
        subgoal_layers: SubgoalLayersInput = None,
        history_subgoals: HistorySubgoalInput | None = None,
        history_max_steps: int | None = None,
    ) -> None:
        """Append one validated root/DAG graph into a caller-owned builder."""
        validate_single_unsupported_lanes(
            actions=actions,
            history_subgoals=history_subgoals,
            history_max_steps=history_max_steps,
        )
        self._runtime.append_into_builder(
            root,
            builder,
            dag=dag,
            goals=goals,
            subgoal_layers=subgoal_layers,
        )

    def encode(
        self,
        root: StateInput,
        dag: TransitionDAG | RXStateDAG | None = None,
        *,
        goals: GoalBatchInput = None,
        actions: ActionBatchInput = None,
        subgoal_layers: SubgoalLayersInput = None,
        history_subgoals: HistorySubgoalInput | None = None,
        history_max_steps: int | None = None,
        include_metadata: bool = True,
        **kwargs,
    ) -> BatchEncoding:
        """Encode one root/DAG pair into native ``BatchEncoding``."""
        return super().encode(
            root,
            goals=goals,
            actions=actions,
            subgoal_layers=subgoal_layers,
            history_subgoals=history_subgoals,
            history_max_steps=history_max_steps,
            dag=dag,
            include_metadata=include_metadata,
            **kwargs,
        )

    def encode_batch(
        self,
        roots: StateBatchInput,
        dags: DagBatchParam = None,
        *,
        goals: GoalBatchParam = None,
        actions: ActionBatchParam = None,
        subgoal_layers: SubgoalLayersBatchParam = None,
        history_subgoals: HistorySubgoalsBatchParam = None,
        history_max_steps: int | None = None,
        batch_attrs: Mapping[str, Any] | None = None,
        collate_spec: CollateSpecParam = None,
        include_metadata: bool = True,
        **kwargs,
    ) -> BatchEncoding:
        """Encode one or many root/DAG pairs into native ``BatchEncoding``."""
        return super().encode_batch(
            roots,
            goals=goals,
            actions=actions,
            subgoal_layers=subgoal_layers,
            history_subgoals=history_subgoals,
            history_max_steps=history_max_steps,
            batch_attrs=batch_attrs,
            collate_spec=collate_spec,
            dags=dags,
            include_metadata=include_metadata,
            **kwargs,
        )

    def _accepted_kwargs(self) -> set[str]:
        """Accept transition DAG kwargs in the generic base API."""
        return super()._accepted_kwargs() | {"dag", "dags"}

    def _encode_batch(
        self,
        roots: StateBatchInput,
        dags: DagBatchParam = None,
        *,
        goals: GoalBatchParam = None,
        subgoal_layers: SubgoalLayersBatchParam = None,
        actions: ActionBatchParam = None,
        history_subgoals: HistorySubgoalsBatchParam = None,
        history_max_steps: int | None = None,
    ) -> BatchEncoding:
        """Encode one or many root/DAG pairs into one batch encoding."""
        validate_batch_unsupported_lanes(
            actions=actions,
            history_subgoals=history_subgoals,
            history_max_steps=history_max_steps,
        )
        return self._runtime.encode_batch(
            roots,
            dags=dags,
            goals=goals,
            subgoal_layers=subgoal_layers,
        )

    def stream(self) -> HorizonEncoderStream:
        """Create a streaming encoder sharing this encoder's C++ engine."""
        return HorizonEncoderStream(self)

    def _horizon_visualization_context(self) -> HorizonVisualizationContext:
        return HorizonVisualizationContext(
            hgraph=self._visualization_context(),
            parent_relation=self.parent_relation,
            sibling_relation=self.sibling_relation,
            cousin_relation=self.cousin_relation,
            target_symbol_prefix=self.target_symbol_prefix,
        )

    def to_networkx(self, data: HeteroData) -> nx.MultiGraph:
        """Convert encoded horizon ``HeteroData`` into a named multigraph."""
        return horizon_to_networkx(data, self._horizon_visualization_context())

    def draw(
        self,
        graph: nx.MultiGraph | HeteroData,
        *,
        ax=None,
        with_labels: bool = True,
        edge_labels: bool = True,
        node_kwargs: dict | None = None,
        edge_kwargs: dict | None = None,
        layout: dict | None = None,
        node_size: float | None = None,
        node_alpha: float | None = None,
        edge_width: float | None = None,
        edge_alpha: float | None = None,
        label_font_size: float | None = None,
        label_nodes: Iterable[str] | None = None,
        label_node_types: Iterable[str] | None = None,
        label_edges: Iterable[tuple[str, ...]] | None = None,
        align_target_nodes: bool = True,
        target_x_spacing: float = 4.0,
        target_y_spacing: float = 2.0,
        layout_seed: int | None = 7,
        symbol_node_scale: float = 1.5,
        non_symbol_linestyle: str | None = "--",
    ):
        if hasattr(graph, "edge_types"):
            graph = self.to_networkx(graph)
        return draw_horizon(
            graph,
            context=self._horizon_visualization_context(),
            ax=ax,
            with_labels=with_labels,
            edge_labels=edge_labels,
            node_kwargs=node_kwargs,
            edge_kwargs=edge_kwargs,
            layout=layout,
            node_size=node_size,
            node_alpha=node_alpha,
            edge_width=edge_width,
            edge_alpha=edge_alpha,
            label_font_size=label_font_size,
            label_nodes=label_nodes,
            label_node_types=label_node_types,
            label_edges=label_edges,
            align_target_nodes=align_target_nodes,
            target_x_spacing=target_x_spacing,
            target_y_spacing=target_y_spacing,
            layout_seed=layout_seed,
            symbol_node_scale=symbol_node_scale,
            non_symbol_linestyle=non_symbol_linestyle,
        )

engine property

Expose the underlying C++ horizon engine.

__init__(domain, *, backend=None, transition_mode=None, target_symbol_prefix=None, target_sources=None, parent_relation=None, sibling_relation=None, cousin_relation=None, enable_parent_relation=None, enable_sibling_relation=None, enable_cousin_relation=None, root_policy=None, max_goal_level=None, symbol_type_id=DEFAULT_SYMBOL_TYPE_ID, ignore_actions=None, add_nullary_predicates=None, include_lgan_edges=None, include_static=None, include_empty_edge_types=None, export_node_names=None, support_literals=None, goal_derivations=None, nullary_object_name=None, lgan_tn_edge_pos=DEFAULT_LGAN_TN_EDGE_POS, lgan_nn_edge_pos=DEFAULT_LGAN_NN_EDGE_POS, lgan_rr_edge_pos=DEFAULT_LGAN_RR_EDGE_POS, history_link_relation=DEFAULT_HISTORY_LINK_RELATION)

Create a hetero horizon encoder.

This lane reads a root state plus a TransitionDAG and creates candidate state rows.

target_sources is therefore simple on this lane:

• state: successor or candidate states from the DAG

root_policy controls how the root state is treated:

• include: encode the root and expose it as a target
• encode_only: encode the root, but omit it from targets and metadata
• exclude: omit the root from targets and keep root facts as base facts

The main-lane sources action, goal, subgoal, and history do not create separate targets here. When include_lgan_edges=True, LGAN anchors are those candidate state rows. There is no lgan_anchor_sources switch here.

Source code in src/mifrost/encoders/horizon.py

def __init__(
    self,
    domain: DomainInput,
    *,
    backend: HorizonBackendName | str | None = None,
    transition_mode: HorizonEncoderMode | str | None = None,
    target_symbol_prefix: str | None = None,
    target_sources: Iterable[TargetSource | str] | None = None,
    parent_relation: str | None = None,
    sibling_relation: str | None = None,
    cousin_relation: str | None = None,
    enable_parent_relation: bool | None = None,
    enable_sibling_relation: bool | None = None,
    enable_cousin_relation: bool | None = None,
    root_policy: RootPolicy | str | None = None,
    max_goal_level: int | None = None,
    symbol_type_id: str | None = DEFAULT_SYMBOL_TYPE_ID,
    ignore_actions: bool | None = None,
    add_nullary_predicates: bool | None = None,
    include_lgan_edges: bool | None = None,
    include_static: bool | None = None,
    include_empty_edge_types: bool | None = None,
    export_node_names: bool | None = None,
    support_literals: bool | None = None,
    goal_derivations: Iterable[Any] | None = None,
    nullary_object_name: str | None = None,
    lgan_tn_edge_pos: str | None = DEFAULT_LGAN_TN_EDGE_POS,
    lgan_nn_edge_pos: str | None = DEFAULT_LGAN_NN_EDGE_POS,
    lgan_rr_edge_pos: str | None = DEFAULT_LGAN_RR_EDGE_POS,
    history_link_relation: str | None = DEFAULT_HISTORY_LINK_RELATION,
) -> None:
    """Create a hetero horizon encoder.

    This lane reads a root state plus a `TransitionDAG` and creates
    candidate state rows.

    `target_sources` is therefore simple on this lane:

    - `state`: successor or candidate states from the DAG

    `root_policy` controls how the root state is treated:

    - `include`: encode the root and expose it as a target
    - `encode_only`: encode the root, but omit it from targets and metadata
    - `exclude`: omit the root from targets and keep root facts as base facts

    The main-lane sources `action`, `goal`, `subgoal`, and `history` do
    not create separate targets here. When `include_lgan_edges=True`, LGAN
    anchors are those candidate state rows. There is no
    `lgan_anchor_sources` switch here.
    """
    normalized_root_policy = normalize_root_policy(root_policy)
    config = self._make_config(
        _HORIZON_CONFIG_CLS,
        symbol_type_id=symbol_type_id,
        target_symbol_prefix=target_symbol_prefix,
        ignore_actions=ignore_actions,
        add_nullary_predicates=add_nullary_predicates,
        include_lgan_edges=include_lgan_edges,
        include_static=include_static,
        include_empty_edge_types=include_empty_edge_types,
        export_node_names=export_node_names,
        max_goal_level=max_goal_level,
        support_literals=support_literals,
        goal_derivations=goal_derivations,
        nullary_object_name=nullary_object_name,
        lgan_tn_edge_pos=lgan_tn_edge_pos,
        lgan_nn_edge_pos=lgan_nn_edge_pos,
        lgan_rr_edge_pos=lgan_rr_edge_pos,
        history_link_relation=history_link_relation,
        transition_mode=_normalize_horizon_mode(transition_mode),
        target_sources=(
            None
            if target_sources is None
            else self._normalize_horizon_target_sources(target_sources)
        ),
        parent_relation=parent_relation,
        sibling_relation=sibling_relation,
        cousin_relation=cousin_relation,
        enable_parent_relation=enable_parent_relation,
        enable_sibling_relation=enable_sibling_relation,
        enable_cousin_relation=enable_cousin_relation,
        root_policy=normalized_root_policy,
    )
    self._runtime = create_horizon_runtime(domain, config, backend=backend)
    self._engine = self._runtime.engine
    self._config = self._engine.config
    self.backend = self._runtime.backend_name
    self.symbol_type_id = self._config.symbol_type_id
    self.target_symbol_prefix = self._config.target_symbol_prefix
    self.parent_relation = self._config.parent_relation
    self.sibling_relation = self._config.sibling_relation
    self.cousin_relation = self._config.cousin_relation
    self.lgan_tn_edge_pos = self._config.lgan_tn_edge_pos
    self.lgan_nn_edge_pos = self._config.lgan_nn_edge_pos
    self.lgan_rr_edge_pos = self._config.lgan_rr_edge_pos
    self.include_lgan_edges = self._config.include_lgan_edges
    self.lgan_anchor_sources = set(self._config.lgan_anchor_sources)
    self._lgan_edge_positions = {
        self.lgan_tn_edge_pos,
        self.lgan_nn_edge_pos,
        self.lgan_rr_edge_pos,
    }

encode(root, dag=None, *, goals=None, actions=None, subgoal_layers=None, history_subgoals=None, history_max_steps=None, include_metadata=True, **kwargs)

Encode one root/DAG pair into native BatchEncoding.

Source code in src/mifrost/encoders/horizon.py

def encode(
    self,
    root: StateInput,
    dag: TransitionDAG | RXStateDAG | None = None,
    *,
    goals: GoalBatchInput = None,
    actions: ActionBatchInput = None,
    subgoal_layers: SubgoalLayersInput = None,
    history_subgoals: HistorySubgoalInput | None = None,
    history_max_steps: int | None = None,
    include_metadata: bool = True,
    **kwargs,
) -> BatchEncoding:
    """Encode one root/DAG pair into native ``BatchEncoding``."""
    return super().encode(
        root,
        goals=goals,
        actions=actions,
        subgoal_layers=subgoal_layers,
        history_subgoals=history_subgoals,
        history_max_steps=history_max_steps,
        dag=dag,
        include_metadata=include_metadata,
        **kwargs,
    )

encode_batch(roots, dags=None, *, goals=None, actions=None, subgoal_layers=None, history_subgoals=None, history_max_steps=None, batch_attrs=None, collate_spec=None, include_metadata=True, **kwargs)

Encode one or many root/DAG pairs into native BatchEncoding.

Source code in src/mifrost/encoders/horizon.py

def encode_batch(
    self,
    roots: StateBatchInput,
    dags: DagBatchParam = None,
    *,
    goals: GoalBatchParam = None,
    actions: ActionBatchParam = None,
    subgoal_layers: SubgoalLayersBatchParam = None,
    history_subgoals: HistorySubgoalsBatchParam = None,
    history_max_steps: int | None = None,
    batch_attrs: Mapping[str, Any] | None = None,
    collate_spec: CollateSpecParam = None,
    include_metadata: bool = True,
    **kwargs,
) -> BatchEncoding:
    """Encode one or many root/DAG pairs into native ``BatchEncoding``."""
    return super().encode_batch(
        roots,
        goals=goals,
        actions=actions,
        subgoal_layers=subgoal_layers,
        history_subgoals=history_subgoals,
        history_max_steps=history_max_steps,
        batch_attrs=batch_attrs,
        collate_spec=collate_spec,
        dags=dags,
        include_metadata=include_metadata,
        **kwargs,
    )

stream()

Create a streaming encoder sharing this encoder's C++ engine.

Source code in src/mifrost/encoders/horizon.py

def stream(self) -> HorizonEncoderStream:
    """Create a streaming encoder sharing this encoder's C++ engine."""
    return HorizonEncoderStream(self)

to_networkx(data)

Convert encoded horizon HeteroData into a named multigraph.

Source code in src/mifrost/encoders/horizon.py

def to_networkx(self, data: HeteroData) -> nx.MultiGraph:
    """Convert encoded horizon ``HeteroData`` into a named multigraph."""
    return horizon_to_networkx(data, self._horizon_visualization_context())