lab_executor ¶

Lab graph executor.

Takes a LiteGraph-serialised graph (from graph.serialize() in the browser) and executes it against a Polars DataFrame, returning the output series keyed by their Signal Out labels.

LiteGraph link format

links: [[link_id, from_node_id, from_slot, to_node_id, to_slot, type], ...]

LabExecutor ¶

Python

LabExecutor(graph: dict[str, Any])

Execute a LiteGraph filter graph against a Polars DataFrame.

Usage::

Text Only
executor = LabExecutor(graph_dict)
outputs  = executor.execute(df)   # {label: pl.Series}

Source code in src/mujoco_mojo/utils/layers/dojo/lab_executor.py

Python
def __init__(self, graph: dict[str, Any]) -> None:
    self._raw_nodes: list[dict] = graph.get("nodes", [])
    self.nodes: dict[int, dict] = {n["id"]: n for n in self._raw_nodes}
    # links keyed by link_id → [link_id, from_node, from_slot, to_node, to_slot, type]
    self.links: dict[int, list] = {lnk[0]: lnk for lnk in graph.get("links", [])}

signal_in_columns `property` ¶

Python

signal_in_columns: list[str]

Column names used by all Signal In nodes - used for validation.

frame_dependencies `property` ¶

Python

frame_dependencies: set[str]

Parquet columns required by any Rotation or Translation node in this graph: the quaternion's x/y/z/w components (Rotation only, when set), the origin's x/y/z components (when set), and the x/y/z siblings of the vector feeding the node's input (which must come directly from a Signal In node - see the rotation/translation handling note in _apply).

output_labels `property` ¶

Python

output_labels: list[str]

Labels produced by all Signal Out nodes.

output_metadata `property` ¶

Python

output_metadata: dict[str, ColumnMetadata]

Metadata declared on each Signal Out node, keyed by output label.

Only the fields ColumnMetadata declares are read (taken from ColumnMetadata.model_fields), so a new field needs no change here. Empty values are ignored, and metadata that fails validation (for example an unknown unit) is skipped with a warning.

execute ¶

Python

execute(
    df: DataFrame,
    template_inputs: dict[str, Series] | None = None,
) -> dict[str, Series]

Run the graph and return {output_label: series}.

Source code in src/mujoco_mojo/utils/layers/dojo/lab_executor.py

Python
def execute(
    self,
    df: pl.DataFrame,
    template_inputs: dict[str, pl.Series] | None = None,
) -> dict[str, pl.Series]:
    """Run the graph and return {output_label: series}."""
    # slot_data[node_id][slot_index] = computed series
    slot_data: dict[int, dict[int, pl.Series]] = defaultdict(dict)

    for nid in self._topo_sort():
        node = self.nodes[nid]
        ntype = self._bare_type(node)
        props = node.get("properties", {})

        if ntype == "signal_in":
            col = props.get("column", "")
            series = (
                df[col].cast(pl.Float64)
                if col in df.columns
                else pl.Series(name=col, values=[0.0] * len(df))
            )
            slot_data[nid][0] = series

        elif ntype == "template_in":
            label = props.get("label", "")
            if template_inputs and label in template_inputs:
                slot_data[nid][0] = template_inputs[label]
            else:
                slot_data[nid][0] = pl.Series(
                    name=label or "_tmpl_in",
                    values=[0.0] * len(df),
                    dtype=pl.Float64,
                )

        elif ntype == "constant":
            value = float(props.get("value", 0.0))
            slot_data[nid][0] = pl.Series(
                name="const", values=[value] * len(df), dtype=pl.Float64
            )

        elif ntype in ("signal_out", "template_out"):
            signal = self._get_input(node, 0, slot_data)
            if signal is not None:
                slot_data[nid][0] = signal

        elif ntype == "template_ref":
            template_graph = props.get("graph")
            if not template_graph:
                continue
            sub = LabExecutor(template_graph)
            # map parent input slots to Template In labels by position
            sub_inputs: dict[str, pl.Series] = {}
            for i, label in enumerate(sub._template_in_labels):
                s = self._get_input(node, i, slot_data)
                if s is not None:
                    sub_inputs[label] = s
            sub_outputs = sub.execute(df, template_inputs=sub_inputs)
            # map Template Out labels to parent output slots by position
            for i, label in enumerate(sub._template_out_labels):
                if label in sub_outputs:
                    slot_data[nid][i] = sub_outputs[label]

        elif ntype == "norm":
            inputs = [
                s.cast(pl.Float64)
                for i in range(len(node.get("inputs", [])))
                if (s := self._get_input(node, i, slot_data)) is not None
            ]
            if not inputs:
                continue
            acc = inputs[0].pow(2)
            for s in inputs[1:]:
                acc = acc + s.pow(2)
            slot_data[nid][0] = acc.sqrt()

        else:
            signal = self._get_input(node, 0, slot_data)
            if signal is None:
                continue
            wrt = self._get_input(node, 1, slot_data)
            slot_data[nid][0] = self._apply(ntype, props, signal, wrt, df)

    # when called as a template sub-execution (template_inputs provided), collect
    # template_out nodes; otherwise collect signal_out nodes for regular lab execution
    terminal = "template_out" if template_inputs is not None else "signal_out"
    outputs: dict[str, pl.Series] = {}
    for nid, node in self.nodes.items():
        if self._bare_type(node) == terminal:
            series = slot_data.get(nid, {}).get(0)
            if series is not None:
                label = node.get("properties", {}).get("label") or f"out_{nid}"
                outputs[label] = series
    return outputs

lab_executor ¶

LabExecutor ¶

signal_in_columns property ¶

frame_dependencies property ¶

output_labels property ¶

output_metadata property ¶

execute ¶

signal_in_columns `property` ¶

frame_dependencies `property` ¶

output_labels `property` ¶

output_metadata `property` ¶