Signal manager

def track(
    self,
    getter: Callable[[], float],
    category: SignalCategory | str,
    subgroups: tuple[str, ...] = (),
    *,
    attr: str | None = None,
    metadata: dict[str, Any] | None = None,
):
    """
    Registers `getter` to be called and posted on every recorded step, under the same `category`/`subgroups`/`attr` namespace as `post`.

    `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
        >>> manager.track(lambda: obj.value, "Custom", ("MyGroup",), attr="value")

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

    """

    def _sample(_: MjState):
        self.post(getter(), category, subgroups, attr=attr, metadata=metadata)

    self.register_sampler(_sample)

def post(
    self,
    value: float,
    category: SignalCategory | str,
    subgroups: tuple[str, ...] = (),
    *,
    attr: str | None = None,
    metadata: dict[str, Any] | None = None,
):
    """
    Injects a value into the telemetry ledger using a hierarchical namespace.

    This method constructs a structured key that the dashboard uses to build a navigable tree view. The naming convention follows a folder-like structure to group related signals (e.g., all axes of a body's position).

    Format:
        Category/Subgroup:Attribute
        (e.g., "Bodies/Link_1:xpos_x")

    Args:
        value (float): The numeric data to record.
        category (SignalCategory | str): Top level category (e.g., "Bodies")
        subgroups (tuple[str, ...], optional): The second-level organizational folders. Defaults to an empty tuple.
        attr (str | None, optional): The specific signal or component name (e.g., "qpos" or "x"). Defaults to None.
        metadata (dict[str, Any] | None, optional): Arbitrary metadata for this signal, persisted into the telemetry file's footer. Only consulted the first time this signal is registered (ignored on later calls for the same signal). Two keys are validated via Pint if present: `dimension` (e.g. `"[length] / [time]"`), for tagging the physical quantity type when the concrete unit isn't knowable (the right choice for built-in signals, since the user's modeling unit system isn't known here), and `unit` (e.g. `"meter / second"`), for the rarer case the concrete unit truly is known. Any other keys (e.g. `display_name`, `comment`) pass through unvalidated. Defaults to None.

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

        >>> # Becomes "Sensors/IMU/Accel:z"
        >>> manager.post(9.81, "Sensors", ("IMU", "Accel"), attr="z")

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

    """
    # use tuple as cache key to avoid string construction
    cache_lookup = (str(category), subgroups, attr if attr is not None else "")

    if cache_lookup in self._key_cache:
        # fast path for cached signal
        full_key = self._key_cache[cache_lookup]
    else:
        # slow path for a new signal
        path_parts = [str(category)] + [str(s) for s in subgroups if s]
        full_key = "/".join(path_parts)
        if attr:
            full_key += f":{attr}"

        self._key_cache[cache_lookup] = full_key

    # get column index
    if full_key in self._key_to_idx:
        idx = self._key_to_idx[full_key]
    else:
        # register a new signal column
        idx = self._n_cols
        self._key_to_idx[full_key] = idx
        self._n_cols += 1

        if metadata:
            self._column_metadata[full_key] = _validate_signal_metadata(metadata)

        logger.debug(f"New signal registered: {full_key} at index {idx}")

        # grow buffer if exceeding the initial guess
        if self._n_cols > self._data_buffer.shape[1]:
            n_cols_to_add = 50
            new_width = self._data_buffer.shape[1] + n_cols_to_add
            logger.debug(f"Growing telemetry buffer width to {new_width} columns.")

            growth = np.full(
                (self._data_buffer.shape[0], n_cols_to_add),
                np.nan,
                dtype=np.float64,
            )
            self._data_buffer = np.hstack([self._data_buffer, growth])

        # more columns means more bytes per row, so the row budget shrinks
        self._recompute_capacity()

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

def record(self, state: MjState):
    """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):
    """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)