"""Notebook utilities.
Small helpers that adapt to the execution context (Jupyter notebook, Colab
or plain Python script): displaying videos, launching tensorboard, and
creating per-run output directories.
"""
import html
import os
import shutil
import sys
import warnings
from datetime import datetime
from pathlib import Path
__all__ = [
"is_notebook",
"video_display",
"setup_tensorboard",
"run_directory",
"outputs_directory",
"silence_known_warnings",
]
def _ipython_shell():
"""The current IPython shell, or `None` when running as a plain script"""
try:
from IPython import get_ipython
return get_ipython()
except ImportError:
return None
[docs]
def is_notebook() -> bool:
"""True when running inside a Jupyter notebook or Google Colab
Used by the helpers that must degrade gracefully when a source is run as a
plain script (as `make check` does).
"""
shell = _ipython_shell()
if shell is None:
return False
return shell.__class__.__name__ == "ZMQInteractiveShell" or "google.colab" in str(
shell.__class__.__module__
)
[docs]
def silence_known_warnings() -> None:
"""Silence third-party warnings that clutter the notebook output
`pygame` (used for rendering) and `tensorboard` both still import the
deprecated `pkg_resources` API: the resulting warning is noise we cannot
fix on our side.
"""
warnings.filterwarnings(
"ignore", message="pkg_resources is deprecated", category=UserWarning
)
[docs]
def outputs_directory() -> Path:
"""The base directory for all outputs (tensorboard logs, videos, ...)
Reads `NOTEBOOK_PROFILE` directly (rather than importing `rl_mind.profiles`,
which needs `cs-lab`) so this stays available with a bare `pip install
rl-mind`; `make check` sets it to `fast-test` (see jupytext-notebook-helper's
`CHECK_PROFILE`).
:returns: `outputs-testing` when `NOTEBOOK_PROFILE` is `fast-test`,
`outputs` otherwise (both relative to the working directory)
"""
profile = os.environ.get("NOTEBOOK_PROFILE", "").strip().upper().replace("-", "_")
return Path("outputs-testing" if profile == "FAST_TEST" else "outputs")
[docs]
def run_directory(name: str) -> Path:
"""Create (and return) a fresh output directory for a training run
:param name: A name for the run, e.g. `"ddpg-CartPoleContinuous-v1-S1"`
:returns: The created directory, `<outputs>/<name>-<YYYYmmdd-HHMMSS>` —
suitable as a tensorboard log directory and as the `run_dir` of an
`Evaluator`
"""
path = outputs_directory() / f"{name}-{datetime.now().strftime('%Y%m%d-%H%M%S')}"
path.mkdir(parents=True, exist_ok=True)
return path
[docs]
def video_display(path: Path | None):
"""Display a video in the notebook (or print its path in a script)
:param path: The video file, or `None` (nothing was recorded)
:returns: An `IPython.display.Video` to be shown as the cell result in a
notebook, `None` otherwise
"""
if path is None:
print("No video to display")
return None
if is_notebook():
from IPython.display import Video
return Video(str(path), embed=True, width=640)
print(f"Video saved in {path}", file=sys.stderr)
return None
[docs]
def setup_tensorboard():
"""Show tensorboard in the notebook, or print the command to launch it
In Jupyter, a chooser lets the student pick where the dashboard lives:
**external** (the default: the command to run from a shell is shown) or
**in the notebook** (the dashboard is displayed inline). Colab has no shell
to run it from, so the dashboard is always shown inline there; from a
plain script (as `make check` runs the sources), the command is printed.
The logs of every run are written under `outputs/` **relative to the
working directory of the kernel** — the absolute path is printed so that a
tensorboard started from a shell watches the right directory.
"""
logdir = outputs_directory().absolute()
logdir.mkdir(parents=True, exist_ok=True)
try:
import tensorboard # noqa: F401
except ImportError:
print(
"The `tensorboard` package is not installed: no visualization "
"will be available (install it with `uv add tensorboard`).",
file=sys.stderr,
)
return
binary = shutil.which("tensorboard") or f"{Path(sys.executable).parent}/tensorboard"
command = f"Launch tensorboard from the shell:\n{binary} --logdir '{logdir}'"
def show_inline() -> bool:
try:
shell = _ipython_shell()
shell.run_line_magic("load_ext", "tensorboard")
shell.run_line_magic("tensorboard", f"--logdir '{logdir}'")
return True
except Exception as exception:
print(f"Could not display tensorboard: {exception}", file=sys.stderr)
return False
if not is_notebook():
print(command)
return
if "google.colab" in str(_ipython_shell().__class__.__module__):
if not show_inline():
print(command)
return
try:
import ipywidgets
from IPython.display import display
except ImportError:
print(command)
return
buttons = ipywidgets.ToggleButtons(
options=[("External", False), ("In the notebook", True)],
value=False,
description="Tensorboard:",
)
# The command goes in a label, not through `print` in an `Output` widget:
# some frontends do not clear an `Output`, and the command would pile up
label = ipywidgets.HTML(f"<pre>{html.escape(command)}</pre>")
output = ipywidgets.Output()
shown = False
def show(inline: bool) -> None:
# The dashboard is loaded once, then only hidden or shown again
nonlocal shown
if inline and not shown:
with output:
shown = show_inline()
label.layout.display = "none" if inline and shown else None
output.layout.display = None if inline else "none"
buttons.observe(lambda change: show(change["new"]), names="value")
display(ipywidgets.VBox([buttons, label, output]))