"""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.
"""
import functools
import io
from typing import Callable, Sequence
import gymnasium as gym
import ipywidgets as widgets
import numpy as np
from IPython.display import display
from PIL import Image, features
from rl_mind.env import VecEnv
__all__ = [
"DecimateTrimmer",
"FrameViewer",
"HistoryEntry",
"RenderWidget",
"Trimmer",
"WindowTrimmer",
"animate_frames",
"decode_frame",
"encode_frame",
"rollout_frames",
]
#: One entry of a frame history: the true step index and the encoded frame
#: (see `encode_frame`)
HistoryEntry = tuple[int, bytes]
@functools.cache
def _frame_format() -> str:
"""WebP when this Pillow build supports it (all the wheels do), else PNG"""
return "webp" if features.check("webp") else "png"
[docs]
def encode_frame(frame: np.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.
:param frame: An `[H, W, 3]` `uint8` array (what `env.render()` returns
with `render_mode="rgb_array"`)
:returns: The encoded image
"""
buffer = io.BytesIO()
options = {"quality": 80, "method": 2} if _frame_format() == "webp" else {}
Image.fromarray(frame).save(buffer, format=_frame_format(), **options)
return buffer.getvalue()
[docs]
def decode_frame(data: bytes) -> np.ndarray:
"""Decode a frame encoded by `encode_frame`
The encoding is lossy: the pixels are close to, not equal to, the
rendered ones.
:param data: The encoded image, e.g. the second element of a
`HistoryEntry`
:returns: An `[H, W, 3]` `uint8` array
"""
return np.asarray(Image.open(io.BytesIO(data)).convert("RGB"))
[docs]
class Trimmer:
"""Bounds how many `(step, frame)` entries a history keeps.
Subclass and override `trim` for custom strategies.
"""
[docs]
def trim(self, history: Sequence[HistoryEntry]) -> Sequence[HistoryEntry]:
"""Return the entries to keep
Called after every new frame, with the whole history (oldest first).
:param history: The current `(step, frame)` entries
:returns: The subsequence to keep (the same object when nothing is
dropped)
"""
raise NotImplementedError
[docs]
class WindowTrimmer(Trimmer):
"""Sliding window: once `history` exceeds `size`, keep only the last
`size` entries (`history[-size:]`, i.e. drop the oldest)."""
def __init__(self, size: int):
""":param size: The number of most recent entries to keep"""
#: The size of the window
self.size = size
[docs]
def trim(self, history: Sequence[HistoryEntry]) -> Sequence[HistoryEntry]:
return history[-self.size :] if len(history) > self.size else history
[docs]
class DecimateTrimmer(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.
"""
def __init__(self, capacity: int, keep_ratio: float = 0.5):
""":param capacity: The number of entries that triggers a trim
:param keep_ratio: The fraction of entries kept at each trim event
(0.5 keeps every other frame)
"""
#: The number of entries that triggers a trim
self.capacity = capacity
#: The subsampling stride, `round(1 / keep_ratio)`
self.stride = max(1, round(1 / keep_ratio))
[docs]
def trim(self, history: Sequence[HistoryEntry]) -> Sequence[HistoryEntry]:
return history[:: self.stride] if len(history) > self.capacity else history
[docs]
class FrameViewer:
"""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.
"""
def __init__(
self, width: str = "400px", fps: int = 30, trimmer: Trimmer | None = None
):
"""Create and display the widget
:param width: The CSS width of the image, e.g. `"400px"`
:param fps: The playback speed of the `Play` button
:param trimmer: Bounds the kept history (see `WindowTrimmer`,
`DecimateTrimmer`); `None` keeps every frame
"""
#: How the history is bounded, if at all
self.trimmer = trimmer
#: The kept `(step, encoded frame)` entries, oldest first
self.history: list[HistoryEntry] = []
#: The number of frames seen since the last `reset`
self.step_count = 0
self.image = widgets.Image(format=_frame_format(), width=width)
self.label = widgets.Label()
self.slider = widgets.IntSlider(min=0, max=0, value=0)
self.play = widgets.Play(min=0, max=0, interval=1000 // fps)
widgets.jslink((self.play, "value"), (self.slider, "value"))
self.slider.observe(self._on_slider, names="value")
display(
widgets.VBox(
[self.image, self.label, widgets.HBox([self.play, self.slider])]
)
)
def _on_slider(self, change) -> None:
step, data = self.history[change["new"]]
self.image.value = data
self.label.value = f"step {step}"
[docs]
def reset(self) -> None:
"""Start a new episode: clears the history"""
self.history = []
self.step_count = 0
[docs]
def update(self, frame: np.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.
:param frame: The rendered frame, an `[H, W, 3]` `uint8` array (what
`env.render()` returns with `render_mode="rgb_array"`)
"""
data = encode_frame(frame)
self.history.append((self.step_count, data))
self.step_count += 1
if self.trimmer is not None:
self.history = list(self.trimmer.trim(self.history))
last = len(self.history) - 1
self.slider.max = last
self.play.max = last
self.slider.value = last # jump to the live frame
self.image.value = data
self.label.value = f"step {self.history[-1][0]}"
[docs]
def rollout_frames(
env_name: str,
policy: Callable[[np.ndarray], object],
seed: int = 0,
trimmer: Trimmer | None = None,
wrappers=(),
**kwargs,
) -> list[HistoryEntry]:
"""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.
:param env_name: The gymnasium environment name
:param policy: a callable `obs -> action`
:param seed: The seed of the episode
:param trimmer: Bounds the returned history (see `WindowTrimmer`,
`DecimateTrimmer`); `None` keeps every frame
:param wrappers: gymnasium wrappers applied to the environment (must match
the ones `policy` was trained with)
:param kwargs: Extra arguments passed to `gymnasium.make`
:returns: a list of `(step, encoded frame)` pairs (see `encode_frame`)
"""
env = gym.make(env_name, render_mode="rgb_array", **kwargs)
for wrapper in wrappers:
env = wrapper(env)
obs, _ = env.reset(seed=seed)
history: list[HistoryEntry] = [(0, encode_frame(env.render()))]
terminated, truncated, step = False, False, 0
while not (terminated or truncated):
obs, reward, terminated, truncated, _ = env.step(policy(obs))
step += 1
history.append((step, encode_frame(env.render())))
if trimmer is not None:
history = list(trimmer.trim(history))
env.close()
return history
[docs]
def animate_frames(
history: Sequence[HistoryEntry], 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`).
:param history: The `(step, encoded frame)` pairs to play, e.g. from
`rollout_frames`
:param fps: The playback speed of the `Play` button
:param width: The CSS width of the image
"""
steps, frames = zip(*history)
image = widgets.Image(value=frames[0], format=_frame_format(), width=width)
label = widgets.Label(value=f"step {steps[0]}")
slider = widgets.IntSlider(min=0, max=len(frames) - 1, value=0)
play = widgets.Play(min=0, max=len(frames) - 1, interval=1000 // fps)
widgets.jslink((play, "value"), (slider, "value"))
def on_change(change):
i = change["new"]
image.value = frames[i]
label.value = f"step {steps[i]}"
slider.observe(on_change, names="value")
display(widgets.VBox([image, label, widgets.HBox([play, slider])]))