signal_manager ¶

def track(self, getter: Callable[[], float], column: Column) -> None:
    """
    Registers `getter` to be called and posted on every recorded step, under `column`.

    `getter` is called fresh on every recorded step, so it should look up a value that changes over the course of the simulation (e.g. a variable updated each step, an attribute, or an indexing operation) rather than a constant computed once. If the underlying value never changes after registration, `track` will simply keep posting that same value every step.

    Examples:
        >>> # Becomes "Custom/MyGroup:value", re-read from `obj.value` every step
        >>> column = Column(category="Custom", subgroups=("MyGroup",), attr="value")
        >>> manager.track(lambda: obj.value, column)

        >>> # Also works: `level` is a variable reassigned each step in the same scope
        >>> level = 0.0
        >>> manager.track(lambda: level, Column(category="Custom", subgroups=("MyGroup",), attr="level"))
        >>> for _ in range(n_steps):
        ...     level = compute_level(state)
        ...     rm.step(state)

    """

    def _sample(_: MjState) -> None:
        self.post(getter(), column)

    self.register_sampler(_sample)

def post(self, value: float, column: Column) -> None:
    """
    Injects a value into the telemetry ledger under `column`.

    The column's name places the signal in a hierarchical namespace that the dashboard uses to build a navigable tree view. Its metadata is persisted into the telemetry file's footer, but only the first time a column with this name is registered (metadata on later columns with the same name is ignored).

    This is called for every column on every timestep, so build each `Column` once and reuse it rather than constructing one per call. A `Column` with the same parts as one already seen still works, but is looked up by comparing its fields instead of by identity.

    Args:
        value (float): The numeric data to record.
        column (Column): Where to record it: the category, subgroups, and attr that make up the column's name, plus its metadata. See `Column` for the naming grammar and `ColumnMetadata` for the metadata fields.

    Examples:
        >>> # Becomes "Bodies/Hand/xpos:x"
        >>> x = Column(category=SignalCategory.BODIES, subgroups=("Hand", "xpos"), attr="x")
        >>> manager.post(1.2, x)

        >>> # Tag a custom signal's physical quantity type without committing to a unit system
        >>> stiffness = Column(
        ...     category="Custom",
        ...     subgroups=("Spring",),
        ...     attr="stiffness",
        ...     metadata={"dimension": "[force] / [length]"},
        ... )
        >>> manager.post(0.4, stiffness)

    """
    idx = self._col_idx.get(column)
    if idx is None:
        # first time this column object is seen: a different column object with the
        # same name may already be registered, in which case it keeps its metadata
        registered = self._key_to_idx.get(str(column))
        idx = registered if registered is not None else self._register(column)
        self._col_idx[column] = idx

    # write value to buffer for next flush
    self._data_buffer[self._buffer_row_idx, idx] = value

def record(self, state: MjState) -> None:
    """Executes all samplers and advances the buffer index. Flushes if due."""
    logger.debug(f"Recording telemetry at t={state.data.time:.6f}")

    self._step_count += 1
    if self._step_count % self.record_decimation != 0:
        return

    # record simulation time
    self._data_buffer[self._buffer_row_idx, 0] = state.data.time

    # run samplers
    for task in self._sample_tasks:
        task(state)

    self._buffer_row_idx += 1

    if self._buffer_row_idx >= self._capacity:
        self.flush()

def flush(self) -> None:
    """Writes the memory buffer to a new part file; parts are merged into `export_path` on `close()`."""
    if self._buffer_row_idx == 0:
        return

    # build column names from mapping
    sorted_keys = sorted(self._key_to_idx.keys(), key=lambda x: self._key_to_idx[x])

    # slice only the used portion of the buffer
    new_df = pl.from_numpy(
        data=self._data_buffer[: self._buffer_row_idx, : self._n_cols],
        schema=sorted_keys,
    )

    part_path = self._part_path(len(self._part_paths))
    logger.info(f"Flushing {self._buffer_row_idx} steps to {part_path.name}")
    # each part is a brand new file, never read back until close()'s merge,
    # so flushing stays O(buffer capacity) instead of O(total rows written
    # so far) and never reads-then-rewrites a file (avoiding a Windows file lock)
    new_df.write_parquet(part_path, compression="zstd")
    self._part_paths.append(part_path)

    # reset buffer for next batch
    self._buffer_row_idx = 0
    self._data_buffer.fill(np.nan)