Skip to content

zymera.viz

viz

zymera.viz — drawing for trajectories and worlds (opt-in extra).

Imports matplotlib/pillow; never imported by the headless core. The viz doctrine is RE-SIMULATION: a run is reproducible from (env spec, checkpoint, key), so render from a fresh rollout(..., keep="all") rather than storing heavy training trajectories.

Deferred (tracked in README status): isometric renderer and keyboard teleop — port from zymera v0 when needed.

draw_frame

draw_frame(ax, world, *, comm_radius=None, show_delivered=True, annotations=())

Draw one world: coverage heat, walls, comm edges, agents (by group).

Source code in zymera/viz/render.py
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
def draw_frame(ax, world, *, comm_radius: Optional[int] = None,
               show_delivered: bool = True, annotations: Sequence = ()) -> None:
    """Draw one world: coverage heat, walls, comm edges, agents (by group)."""
    h, w = world.grid_h, world.grid_w
    explored = np.asarray(world.explored, dtype=float)
    covered = np.asarray(world.covered, dtype=float)
    wall = np.asarray(world.wall)
    pos = np.asarray(world.body.position)
    group = np.asarray(world.group)

    base = 0.25 * covered + 0.06 * np.minimum(explored, 8)
    base[wall] = np.nan                                   # walls drawn as ink
    ax.imshow(base, cmap="YlGn", vmin=0, vmax=1, origin="upper")
    ax.imshow(np.where(wall, 1.0, np.nan), cmap="gray_r", vmin=0, vmax=1,
              origin="upper")

    if comm_radius is not None:                           # potential: thin dotted
        for i, j in _comm_edges(pos, comm_radius):
            ax.plot([pos[i, 1], pos[j, 1]], [pos[i, 0], pos[j, 0]],
                    ls=":", lw=0.9, color="0.45", zorder=2)
    if show_delivered:                                    # delivered: solid
        adj = np.asarray(world.comm_graph)
        n = adj.shape[0]
        for i in range(n):
            for j in range(i + 1, n):
                if adj[i, j]:
                    ax.plot([pos[i, 1], pos[j, 1]], [pos[i, 0], pos[j, 0]],
                            lw=1.6, color="#2e6e63", zorder=3)

    cmap = plt.get_cmap("tab10")
    multi_group = len(np.unique(group)) > 1
    colors = [cmap(int(group[i]) if multi_group else i % 10)
              for i in range(pos.shape[0])]
    ax.scatter(pos[:, 1], pos[:, 0], c=colors, s=120, edgecolors="black",
               linewidths=0.8, zorder=4)

    for a in annotations:                                 # mission overlays
        if isinstance(a, Point):
            p = np.asarray(a.pos)
            ax.scatter([p[1]], [p[0]], marker="*", s=180, color="#8e2c1f",
                       zorder=5)
        elif isinstance(a, PathAnn):
            cells = np.asarray(a.cells)
            ax.plot(cells[:, 1], cells[:, 0], lw=1.4, color="#8e2c1f",
                    alpha=0.8, zorder=5)
        elif isinstance(a, Region):
            mask = np.asarray(a.mask, dtype=float)
            ax.imshow(np.where(mask > 0, 0.8, np.nan), cmap="Reds", vmin=0,
                      vmax=1, alpha=0.3, origin="upper")

    ax.set_xticks([]), ax.set_yticks([])
    ax.set_xlim(-0.5, w - 0.5), ax.set_ylim(h - 0.5, -0.5)

render_comm_gif

render_comm_gif(worlds, path, *, comm_radius, fps=6, annotations=None)

GIF with the comm overlay: potential edges dotted, delivered solid.

Source code in zymera/viz/render.py
115
116
117
118
119
def render_comm_gif(worlds, path: str, *, comm_radius: int, fps: int = 6,
                    annotations=None) -> str:
    """GIF with the comm overlay: potential edges dotted, delivered solid."""
    return render_gif(worlds, path, fps=fps, comm_radius=comm_radius,
                      show_delivered=True, annotations=annotations)

render_gif

render_gif(worlds, path, *, fps=6, comm_radius=None, show_delivered=True, annotations=None)

Trajectory -> animated GIF at path. Returns the path.

Source code in zymera/viz/render.py
105
106
107
108
109
110
111
112
def render_gif(worlds, path: str, *, fps: int = 6, comm_radius=None,
               show_delivered: bool = True, annotations=None) -> str:
    """Trajectory -> animated GIF at ``path``. Returns the path."""
    frames = render_frames(worlds, comm_radius=comm_radius,
                           show_delivered=show_delivered, annotations=annotations)
    frames[0].save(path, save_all=True, append_images=frames[1:],
                   duration=int(1000 / fps), loop=0)
    return path

make_report

make_report(traj, path, *, env=None, title='zymera report', fps=6, comm_radius=None)

Write a single-file HTML report for a rollout(..., keep="all") dict.

Sections: episode GIF, coverage-over-time, per-term reward curves (when the rollout used collect=("reward_terms",)), final visit heatmap, and the env composition when env is given.

Source code in zymera/viz/report.py
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
def make_report(traj, path: str, *, env=None, title: str = "zymera report",
                fps: int = 6, comm_radius: Optional[int] = None) -> str:
    """Write a single-file HTML report for a ``rollout(..., keep="all")`` dict.

    Sections: episode GIF, coverage-over-time, per-term reward curves (when the
    rollout used ``collect=("reward_terms",)``), final visit heatmap, and the
    env composition when ``env`` is given.
    """
    worlds = traj["world"]
    if comm_radius is None and env is not None:
        topo = getattr(getattr(env, "channel", None), "topology", None)
        comm_radius = getattr(topo, "radius", None)

    frames = render_frames(worlds, comm_radius=comm_radius)
    gif64 = _b64_gif(frames, fps)

    covered = np.asarray(worlds.seen_by).any(1)                # (T+1, H, W)
    cov_curve = covered.reshape(covered.shape[0], -1).mean(-1)
    fig, ax = plt.subplots(figsize=(5, 2.4))
    ax.plot(cov_curve, color="#8e2c1f")
    ax.set_xlabel("step"), ax.set_ylabel("coverage"), ax.set_ylim(0, 1)
    ax.grid(alpha=0.25)
    cov64 = _b64_fig(fig)

    terms_html = ""
    terms = traj.get("info", {}).get("reward_terms")
    if terms:
        fig, ax = plt.subplots(figsize=(5, 2.4))
        for name, arr in sorted(terms.items()):
            ax.plot(np.asarray(arr).mean(-1), label=name, lw=1.2)
        ax.set_xlabel("step"), ax.set_ylabel("unweighted term mean")
        ax.legend(fontsize=7), ax.grid(alpha=0.25)
        terms_html = (
            "<h2>Reward terms</h2>"
            f'<img alt="per-term reward curves" src="data:image/png;base64,{_b64_fig(fig)}">'
        )

    fig, ax = plt.subplots(figsize=(3.2, 3.2))
    ax.imshow(np.asarray(worlds.explored)[-1], cmap="YlGn", origin="upper")
    ax.set_xticks([]), ax.set_yticks([]), ax.set_title("visit counts", fontsize=9)
    heat64 = _b64_fig(fig)

    env_html = ""
    if env is not None:
        env_html = f"<h2>Env composition</h2><pre>{html.escape(repr(env))}</pre>"

    doc = f"""<!DOCTYPE html>
<html><head><meta charset="utf-8"><title>{html.escape(title)}</title>
<style>
 body {{ font-family: Charter, 'Iowan Old Style', Georgia, serif; margin: 2rem auto;
        max-width: 46rem; color: #1a1816; background: #f7f4ee; }}
 h1, h2 {{ font-family: 'Avenir Next', Avenir, 'Helvetica Neue', sans-serif; }}
 h1 {{ border-bottom: 2px solid #8e2c1f; padding-bottom: .3rem; }}
 img {{ max-width: 100%; border: 1px solid #d8d2c4; border-radius: 6px; }}
 pre {{ background: #211e1b; color: #f0ece2; padding: .8rem 1rem; border-radius: 6px;
       font: 12px 'SF Mono', Menlo, monospace; overflow-x: auto; }}
</style></head><body>
<h1>{html.escape(title)}</h1>
<h2>Episode</h2><img alt="episode gif" src="data:image/gif;base64,{gif64}">
<h2>Coverage over time</h2><img alt="coverage curve" src="data:image/png;base64,{cov64}">
{terms_html}
<h2>Final visit heatmap</h2><img alt="visit heatmap" src="data:image/png;base64,{heat64}">
{env_html}
</body></html>"""
    with open(path, "w") as f:
        f.write(doc)
    return path