Rendering — rl_mind.rendering

Live rendering of environment frames in a Jupyter widget.

RenderWidget wraps an env (created with render_mode="rgb_array") so every reset/step renders a frame and pushes it to a live-updating ipywidgets.Image, with a slider and a Play button to scrub back through what happened. It accepts either a single gymnasium env:

env = RenderWidget(gym.make("CartPole-v1", render_mode="rgb_array"))
obs, _ = env.reset(seed=0)
obs, reward, terminated, truncated, _ = env.step(action)

or a VecEnv — which steps every copy together, so k selects which one to render (via VecEnv.gym_unwrapped, the existing escape hatch):

env = RenderWidget(
    VecEnv("CartPole-v1", num_envs=4, seed=0, render_mode="rgb_array"), k=1
)
obs = env.reset()
step = env.step(actions)

For a frame source with neither API (e.g. driving the widget from something that isn’t stepped as an env at all), use the underlying FrameViewer directly and call FrameViewer.update / FrameViewer.reset yourself.

Trimmer bounds how much history a FrameViewer (or rollout_frames) keeps in memory: WindowTrimmer is a sliding window over the most recent steps, DecimateTrimmer instead subsamples so a long episode stays bounded in memory while still spanning the whole run (older parts get sparser).

rollout_frames/animate_frames cover the complementary case: play back a fixed episode (e.g. already collected, or from a trained actor) rather than following one live.

Histories keep frames encoded (lossy WebP, see encode_frame): a raw MuJoCo frame weighs ~700 KB, its WebP ~10 KB, so a whole 1000-step episode fits in memory without a trimmer. decode_frame gives the pixels back.

rl_mind.rendering.HistoryEntry

One entry of a frame history: the true step index and the encoded frame (see encode_frame)

alias of tuple[int, bytes]

Watching an environment live:

class rl_mind.rendering.RenderWidget(env: Env | VecEnv, width: str = '400px', fps: int = 30, trimmer: Trimmer | None = None, k: int = 0)

Bases: object

Wrap an env supporting render_mode="rgb_array" in a FrameViewer: every reset/step renders and appends a frame automatically.

env can be a single gym.Env, or a VecEnv (which steps every copy together and has no single-environment gym API of its own): k then selects which copy to render, via VecEnv.gym_unwrapped. Attribute access (e.g. action_space, n_actions) is forwarded to env.

Parameters:

• env – The environment to wrap, created with render_mode="rgb_array"

• width – The CSS width of the image, e.g. "400px"

• fps – The playback speed of the Play button

• trimmer – bounds/subsamples the kept history (see WindowTrimmer, DecimateTrimmer); None keeps everything

• k – index of the copy to render, when env is a VecEnv

env

The wrapped environment

k

Which copy of a VecEnv is rendered

viewer

The widget the frames are pushed to

reset(*args, **kwargs)

Reset the wrapped environment, clearing the frame history

Parameters:

args – Passed to the wrapped reset (e.g. seed=0 for a gym.Env; VecEnv.reset takes no argument)

Returns:

Whatever the wrapped environment returns

step(action)

Step the wrapped environment and render the resulting frame

Parameters:

action – Passed to the wrapped step (a single action for a gym.Env, a [B, ...] tensor of actions for a VecEnv)

Returns:

Whatever the wrapped environment returns (a gymnasium 5-tuple, or an rl_mind.env.EnvStep for a VecEnv)

class rl_mind.rendering.FrameViewer(width: str = '400px', fps: int = 30, trimmer: Trimmer | None = None)

Bases: object

A live-updating Jupyter widget for a stream of rendered frames.

Call update(frame) each time a new frame is available (e.g. after every environment step) and reset() at the start of a new episode. A slider and a Play button let you scrub back through the (optionally trimmed) history; a label shows the true step number even when the history is sparse.

trimmer

How the history is bounded, if at all

history: list[tuple[int, bytes]]

The kept (step, encoded frame) entries, oldest first

step_count

The number of frames seen since the last reset

reset() → None

Start a new episode: clears the history

update(frame: ndarray) → None

Append a newly rendered frame, advancing the step counter

The slider jumps to the new frame, so the widget follows what is happening live; scrubbing back is always possible afterwards.

Parameters:

frame – The rendered frame, an [H, W, 3] uint8 array (what env.render() returns with render_mode="rgb_array")

Bounding how much history is kept:

class rl_mind.rendering.Trimmer

Bases: object

Bounds how many (step, frame) entries a history keeps.

Subclass and override trim for custom strategies.

trim(history: Sequence[tuple[int, bytes]]) → Sequence[tuple[int, bytes]]

Return the entries to keep

Called after every new frame, with the whole history (oldest first).

Parameters:

history – The current (step, frame) entries

Returns:

The subsequence to keep (the same object when nothing is dropped)

class rl_mind.rendering.WindowTrimmer(size: int)

Bases: Trimmer

Sliding window: once history exceeds size, keep only the last size entries (history[-size:], i.e. drop the oldest).

size

The size of the window

trim(history: Sequence[tuple[int, bytes]]) → Sequence[tuple[int, bytes]]

Return the entries to keep

Called after every new frame, with the whole history (oldest first).

Parameters:

history – The current (step, frame) entries

Returns:

The subsequence to keep (the same object when nothing is dropped)

class rl_mind.rendering.DecimateTrimmer(capacity: int, keep_ratio: float = 0.5)

Bases: Trimmer

Bounded memory (capacity), but spans the whole history: once the cap is hit, subsample with a stride so the buffer drops to roughly capacity * keep_ratio entries.

keep_ratio is the fraction kept at each trim event, not of the whole episode: since capacity is bounded but episode length isn’t, older segments get sparser the longer things run — there is no way around that with O(1) memory.

capacity

The number of entries that triggers a trim

stride

The subsampling stride, round(1 / keep_ratio)

trim(history: Sequence[tuple[int, bytes]]) → Sequence[tuple[int, bytes]]

Return the entries to keep

Called after every new frame, with the whole history (oldest first).

Parameters:

history – The current (step, frame) entries

Returns:

The subsequence to keep (the same object when nothing is dropped)

Replaying a fixed episode:

rl_mind.rendering.rollout_frames(env_name: str, policy: Callable[[ndarray], object], seed: int = 0, trimmer: Trimmer | None = None, wrappers=(), **kwargs) → list[tuple[int, bytes]]

Step through one episode with policy (obs -> action) on any env supporting render_mode="rgb_array", collecting (step, encoded frame) pairs (optionally bounded/subsampled by trimmer); step is the true step index, preserved even when trimmer makes the history sparse.

Parameters:

• env_name – The gymnasium environment name

• policy – a callable obs -> action

• seed – The seed of the episode

• trimmer – Bounds the returned history (see WindowTrimmer, DecimateTrimmer); None keeps every frame

• wrappers – gymnasium wrappers applied to the environment (must match the ones policy was trained with)

• kwargs – Extra arguments passed to gymnasium.make

Returns:

a list of (step, encoded frame) pairs (see encode_frame)

rl_mind.rendering.animate_frames(history: Sequence[tuple[int, bytes]], fps: int = 30, width: str = '400px') → None

Play a (step, encoded frame) history with play/pause, step and rewind controls; the label shows the true step number even when history is sparse (e.g. produced with a Trimmer).

Parameters:

• history – The (step, encoded frame) pairs to play, e.g. from rollout_frames

• fps – The playback speed of the Play button

• width – The CSS width of the image

rl_mind.rendering.encode_frame(frame: ndarray) → bytes

Encode a rendered frame, as kept in histories and sent to the widgets

Lossy WebP (quality 80, method=2): ~5 ms and ~10 KB for a MuJoCo frame, against ~700 KB raw — slower settings or lossless encoding cost up to ten times as much for a barely smaller file. Falls back to PNG if this Pillow build lacks WebP.

Parameters:

frame – An [H, W, 3] uint8 array (what env.render() returns with render_mode="rgb_array")

Returns:

The encoded image

rl_mind.rendering.decode_frame(data: bytes) → ndarray

Decode a frame encoded by encode_frame

The encoding is lossy: the pixels are close to, not equal to, the rendered ones.

Parameters:

data – The encoded image, e.g. the second element of a HistoryEntry

Returns:

An [H, W, 3] uint8 array