torch_world.plots

Matplotlib figures for world models; requires pip install torch-world[examples].

One visual system across figures: a fixed categorical order, one sequential hue for time, hairline recessive axes, 2px lines, ringed markers and ink-colored text. In notebooks, figures render at retina resolution.

style

style()

images

images(
    x: Tensor,
    titles: list[str] | None = None,
    ncols: int = 8,
    size: float = 1.2,
    rows: list[str] | None = None,
)

Grid of (N, 3, H, W) images (float in [0, 1] or uint8), with optional titles and rows labels.

masks

masks(
    image: Tensor,
    grid: tuple[int, int],
    layers: dict[str, Tensor],
    size: float = 2.0,
    color: str | None = None,
    title: str | None = None,
)

The image beside one panel per mask, tinting the patches each index set holds on grid.

Each mask takes the next categorical color, or color for all of them when they show one mask over time.

curves

curves(
    series: dict,
    smooth: int = 25,
    logy: bool = True,
    ylabel: str = "loss",
    size=(5.2, 2.6),
)

Training curves: the raw values as a faint trace under a smooth-point moving average, labeled at the end.

Each series is a list of values, one per step, or a (steps, values) pair when it is logged on some steps only.

trajectories

trajectories(
    paths: Tensor,
    goals: Tensor,
    radius: float = 0.1,
    ncols: int = 4,
    size: float = 1.9,
    background: Tensor | None = None,
    scale: float = 1.0,
    success: Tensor | None = None,
)

Small multiples of planar paths (N, T, 2) in the unit square toward goals (N, 2).

The path darkens with time on one sequential hue; the start is an open ring and the end a filled diamond, so both stay visible when they coincide, and the goal is a star inside its success disc of radius. Each title marks the episode a success (✓) or failure (✗) and reports the distance to the goal at the start and at the end. success (N,) overrides the default test, ending inside the disc. Titles and the legend show distances times scale, for example meters per unit of the square.

distances

distances(
    runs: dict[str, Tensor],
    radius: float | None = None,
    size=(5.2, 2.6),
)

Distance to goal per step (N, T) for each method: the median as a line over the interquartile band.

bars

bars(
    groups: dict[str, dict[str, float]],
    xlabel: str = "accuracy",
    baseline: str | None = None,
    size=None,
)

Horizontal grouped bars: one row per key of groups, one bar per series, each labeled with its value.

Series take the categorical colors in order; the baseline series is drawn in neutral gray. Values in [0, 1] are shown as percentages. Values with low and high (metrics.Rate, e.g. from GoalReaching.intervals()) get a whisker for their confidence interval.

answers

answers(
    x: Tensor,
    questions: list[str],
    predictions: list[str],
    truths: list[str],
    ncols: int = 4,
    size: float = 2.1,
)

Question-answer cards: each image under its question, with the predicted answer marked right or wrong.

probes

probes(
    true: Tensor,
    pred: Tensor,
    names: list[str],
    size: float = 2.5,
)

Linear-probe scatter per target (N, K): prediction against truth on the identity line, with R² in the title.