cpomdp-workshop 0.1.0__py3-none-any.whl

cpomdp_workshop/__init__.py

@@ -0,0 +1,161 @@
+"""Build a curious rocket: the workshop package.
+
+Everything the notebook imports lives here. The world and its sensor field are in
+``world``, the compass moves and the search over them in ``planner``, the two rockets
+in ``build``, the flight loop in ``sim`` and the checks the room runs in ``check``.
+
+The cpomdp release the workshop was measured against is pinned in ``pyproject.toml``
+and read back here as ``WORKSHOP_CPOMDP_VERSION``, so the notebook, the tests and the
+install cannot disagree about which library produced a number.
+"""
+
+import sys
+from importlib import metadata
+
+from cpomdp_workshop.build import (
+    THRUSTER_COST,
+    curious_rocket,
+    model_with,
+    one_node,
+    pad_goal,
+    pad_preference,
+    the_1960_rocket,
+)
+from cpomdp_workshop.check import (
+    check_agent,
+    check_backends,
+    check_my_sensor,
+    epistemic_spread,
+    grade_sensor,
+    probe,
+    show_probe,
+)
+from cpomdp_workshop.checkpoints import (
+    after_session_1,
+    after_session_2,
+    after_session_3,
+)
+from cpomdp_workshop.planner import (
+    STEP,
+    compass_actions,
+    compass_moves,
+    compass_planner,
+)
+from cpomdp_workshop.plot import (
+    plot_beliefs,
+    plot_flights,
+    plot_path,
+    plot_world,
+    preview,
+    world_flights,
+)
+from cpomdp_workshop.selfguided import two_node_backend
+from cpomdp_workshop.sim import ROCK, Flight, fly
+from cpomdp_workshop.solutions import (
+    CLEAR_SKY_ROCK,
+    axis_certificate,
+    efe_split,
+    rxinfer_refusal,
+)
+from cpomdp_workshop.symptoms import symptoms
+from cpomdp_workshop.world import (
+    BEACON,
+    BLACKOUT,
+    LAUNCH,
+    PAD,
+    ROCK_RADIUS,
+    beacon_sensor,
+    default_params,
+    reference_sensor,
+)
+
+__all__ = [
+    "BEACON",
+    "BLACKOUT",
+    "CLEAR_SKY_ROCK",
+    "LAUNCH",
+    "MINIMUM_PYTHON",
+    "PAD",
+    "ROCK",
+    "ROCK_RADIUS",
+    "STEP",
+    "THRUSTER_COST",
+    "WORKSHOP_CPOMDP_VERSION",
+    "Flight",
+    "after_session_1",
+    "after_session_2",
+    "after_session_3",
+    "axis_certificate",
+    "beacon_sensor",
+    "check_agent",
+    "check_backends",
+    "check_my_sensor",
+    "compass_actions",
+    "compass_moves",
+    "compass_planner",
+    "curious_rocket",
+    "default_params",
+    "efe_split",
+    "epistemic_spread",
+    "fly",
+    "grade_sensor",
+    "hello",
+    "model_with",
+    "one_node",
+    "pad_goal",
+    "pad_preference",
+    "plot_beliefs",
+    "plot_flights",
+    "plot_path",
+    "plot_world",
+    "preview",
+    "probe",
+    "reference_sensor",
+    "rxinfer_refusal",
+    "show_probe",
+    "symptoms",
+    "the_1960_rocket",
+    "two_node_backend",
+    "world_flights",
+]
+
+MINIMUM_PYTHON = (3, 11)
+
+
+def _pinned_cpomdp_version() -> str:
+    """The cpomdp version this distribution pins, read from its own metadata."""
+    for requirement in metadata.requires("cpomdp-workshop") or ():
+        name, pinned, rest = requirement.partition("==")
+        if name.strip() == "cpomdp" and pinned:
+            return rest.split(";")[0].strip()
+    raise RuntimeError(
+        "cpomdp-workshop must pin cpomdp with '==' in pyproject.toml. The notebook "
+        "and the tests read the pin from there."
+    )
+
+
+WORKSHOP_CPOMDP_VERSION = _pinned_cpomdp_version()
+
+
+def hello() -> None:
+    """Print the versions in use and the world, after checking the install.
+
+    Raises:
+        RuntimeError: if Python is older than 3.11 or the installed cpomdp is not
+            the release the workshop was measured on.
+    """
+    python = ".".join(str(part) for part in sys.version_info[:3])
+    if sys.version_info[:2] < MINIMUM_PYTHON:
+        raise RuntimeError(f"Python {python} is too old. The workshop needs 3.11+.")
+    installed = metadata.version("cpomdp")
+    if installed != WORKSHOP_CPOMDP_VERSION:
+        raise RuntimeError(
+            f"cpomdp {installed} is installed and the workshop was measured on "
+            f"{WORKSHOP_CPOMDP_VERSION}. Run: pip install cpomdp=="
+            f"{WORKSHOP_CPOMDP_VERSION}"
+        )
+    print(f"cpomdp {installed}, jax {metadata.version('jax')}, Python {python}. Ready.")
+    print(
+        f"launch {LAUNCH} -> pad {PAD}, nav beacon at {BEACON}, "
+        f"blackout at {BLACKOUT} with a rock of radius {ROCK_RADIUS}"
+    )

cpomdp_workshop/_seams.py

@@ -0,0 +1,46 @@
+"""Where this package meets cpomdp past its top-level exports.
+
+The finite-action search, the per-step EFE trace and the RxInfer oracle live in
+submodules the library does not re-export at its top level. They are imported here and
+nowhere else, so a rename in one of them is a fix to this file alone. The notebook
+reaches them through the package and never imports a submodule itself.
+
+The one typing workaround the pinned release needs lives here too.
+"""
+
+from typing import cast
+
+from cpomdp import CallableSensor, FixedSensor, ObservationModel
+from cpomdp.backends.rxinfer import RxInferBackend
+from cpomdp.efe import PolicyEFETrace, policy_efe_trace
+from cpomdp.enumeration import (
+    ChunkedEFESearch,
+    CompletenessCertificate,
+    EnumeratedEFESearch,
+    ExhaustiveEFESearch,
+    FiniteActionSet,
+    RecedingHorizonSelector,
+)
+
+__all__ = [
+    "ChunkedEFESearch",
+    "CompletenessCertificate",
+    "EnumeratedEFESearch",
+    "ExhaustiveEFESearch",
+    "FiniteActionSet",
+    "PolicyEFETrace",
+    "RecedingHorizonSelector",
+    "RxInferBackend",
+    "as_observation_model",
+    "policy_efe_trace",
+]
+
+
+def as_observation_model(sensor: FixedSensor | CallableSensor) -> ObservationModel:
+    """``sensor``, typed as the protocol it implements.
+
+    cpomdp 0.5.0 declares the protocol's ``is_fixed`` writable and sets it as a class
+    attribute on both sensors, so a type checker refuses either sensor where the
+    protocol is expected. The fix belongs in the library.
+    """
+    return cast(ObservationModel, sensor)

cpomdp_workshop/build.py

@@ -0,0 +1,171 @@
+"""The two rockets, and the 1960 one again as a graph of one node.
+
+Both rockets fly the same plant: a position in the plane, a move added straight to it,
+a reading of where they are. They differ in the reading. The 1960 rocket's reading is
+equally noisy everywhere, so no move can teach it more than another. It is handed a
+state to reach and steers by LQR. The curious rocket's reading depends on where it is.
+It is handed a reading to prefer and searches the compass moves for the plan that
+best trades getting there against seeing clearly.
+"""
+
+import numpy as np
+from cpomdp import (
+    Agent,
+    Belief,
+    CallableSensor,
+    CouplingGraph,
+    CouplingGraphBackend,
+    FixedSensor,
+    GaussianTransition,
+    LinearGaussianModel,
+    Preference,
+    StateGoal,
+)
+
+from cpomdp_workshop._seams import as_observation_model
+from cpomdp_workshop.planner import compass_planner
+from cpomdp_workshop.world import (
+    DYNAMICS_NOISE,
+    LAUNCH,
+    NOISE_FLOOR,
+    PAD,
+    PRIOR_COV,
+    NoiseFn,
+    SensorParams,
+    default_params,
+    reference_sensor,
+)
+
+__all__ = [
+    "THRUSTER_COST",
+    "curious_rocket",
+    "model_with",
+    "one_node",
+    "pad_goal",
+    "pad_preference",
+    "the_1960_rocket",
+]
+
+THRUSTER_COST = 100.0  # effort
+
+_IDENTITY = np.eye(2)
+
+
+def model_with(
+    sensor: NoiseFn | None = None, params: SensorParams | None = None
+) -> LinearGaussianModel:
+    """The rocket's generative model, with or without a state-dependent sensor.
+
+    Args:
+        sensor: a noise function ``R(x, params)``, or ``None`` for a reading that is
+            equally noisy everywhere.
+        params: the sensor's parameters. Required when ``sensor`` is given.
+
+    Returns:
+        A two-dimensional model. The rocket stays where it is unless moved, a move
+        adds straight to the position, and the reading is the position plus noise.
+
+    Raises:
+        ValueError: if a sensor is given without parameters.
+    """
+    if sensor is None:
+        observation_model = None
+        observation_matrix, observation_noise = (
+            _IDENTITY,
+            NOISE_FLOOR * _IDENTITY,
+        )  # C, R
+    elif params is None:
+        raise ValueError("a sensor needs its params. Pass default_params(k).")
+    else:
+        # The model reads C, and R at the prior mean, from the sensor itself.
+        observation_model = as_observation_model(
+            CallableSensor(_IDENTITY, sensor, params)
+        )
+        observation_matrix = observation_noise = None
+    return LinearGaussianModel(
+        _IDENTITY,
+        control_matrix=_IDENTITY,
+        observation_matrix=observation_matrix,
+        observation_noise=observation_noise,
+        dynamics_noise=DYNAMICS_NOISE * _IDENTITY,
+        prior=Belief(mean=np.asarray(LAUNCH), cov=PRIOR_COV * _IDENTITY),
+        observation_model=observation_model,
+    )
+
+
+def pad_goal() -> StateGoal:
+    """The pad as a state to reach, with every unit of thrust charged for.
+
+    The effort weight is the thruster cost. Left at its default of one, LQR would
+    cover more than half the distance to the pad in its first move.
+    """
+    return StateGoal(
+        np.asarray(PAD), precision=_IDENTITY, effort=THRUSTER_COST * _IDENTITY
+    )
+
+
+def pad_preference() -> Preference:
+    """The pad as a reading to prefer, which is what a compass search scores."""
+    return Preference(np.asarray(PAD), precision=_IDENTITY)
+
+
+def the_1960_rocket() -> Agent:
+    """The Kalman filter with a goal bolted on: one fixed reading, steered by LQR."""
+    return Agent(model_with(), pad_goal())
+
+
+def curious_rocket(
+    beacon_gain: float, sensor: NoiseFn = reference_sensor, *, horizon: int = 1
+) -> Agent:
+    """A rocket whose reading depends on where it is, searching the compass moves.
+
+    Args:
+        beacon_gain: k, how fast the reading degrades away from the beacon.
+        sensor: the noise field the rocket models. The default has the beacon and
+            the blackout. ``beacon_sensor`` is the beacon alone.
+        horizon: how many moves ahead the search looks, H.
+
+    Returns:
+        An agent that prefers to read the pad and plans every step over every
+        compass plan of length ``horizon``.
+    """
+    model = model_with(sensor, default_params(beacon_gain))
+    planner = compass_planner(model, horizon=horizon)
+    return Agent(model, pad_preference(), selector=planner)
+
+
+def one_node(model: LinearGaussianModel) -> CouplingGraphBackend:
+    """The same model as a graph of one node, to run it on the message-passing engine.
+
+    A coupling graph is for branching models, where separate variables evolve at
+    their own rates and only some of them are read. A single node with no edges is
+    the degenerate case: the filtering problem the Kalman backend solves, routed
+    through the graph code instead. The two engines should agree to the last few
+    digits.
+
+    Args:
+        model: a fixed-sensor model from [`model_with`][cpomdp_workshop.model_with].
+
+    Returns:
+        A backend for ``Agent.from_backend``. Its graph is ``backend.graph``.
+
+    Raises:
+        ValueError: if the model's sensor depends on the state. The backends would
+            then not be running the same model.
+    """
+    if model.observation_model is not None:
+        raise ValueError(
+            "one_node takes a fixed-sensor model. Build it with model_with()."
+        )
+    sensor = as_observation_model(
+        FixedSensor(model.observation_matrix, observation_noise=model.observation_noise)
+    )
+    graph = CouplingGraph(
+        dims=[model.n_states], couplings=[], observation_models={0: sensor}
+    )
+    transition = GaussianTransition(
+        model.dynamics_matrix, dynamics_noise=model.dynamics_noise
+    )
+    return CouplingGraphBackend(
+        graph, [transition], prior=model.prior, control_matrix=model.control_matrix
+    )

cpomdp_workshop/check.py

@@ -0,0 +1,287 @@
+"""What the notebook's checks compute, and what they print.
+
+Each check is a pure function the tests assert on, paired with a ``check_`` or
+``show_`` function that prints its result. A sensor is graded by flying a rocket with
+it, never passed or failed. One that varies and still leaves the path straight is
+told which way to move k.
+"""
+
+from dataclasses import dataclass
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+from cpomdp import Agent, InferenceBackend, LinearGaussianModel, ProbeReport
+from cpomdp import probe_model as _probe_model
+
+from cpomdp_workshop.build import model_with, pad_goal, pad_preference
+from cpomdp_workshop.planner import compass_moves, compass_planner
+from cpomdp_workshop.sim import fly
+from cpomdp_workshop.world import (
+    LAUNCH,
+    PAD,
+    NoiseFn,
+    SensorParams,
+    beacon_sensor,
+    reference_sensor,
+)
+
+__all__ = [
+    "LANDED",
+    "STRAIGHT_BEND",
+    "SensorGrade",
+    "backend_differences",
+    "check_agent",
+    "check_backends",
+    "check_my_sensor",
+    "epistemic_spread",
+    "grade_sensor",
+    "probe",
+    "show_probe",
+]
+
+STRAIGHT_BEND = 0.1
+"""A path that strays less than this from the launch-to-pad line counts as straight."""
+LANDED = 0.5
+"""A flight that ends further than this from the pad has not landed."""
+
+_ROUTE = np.linspace(LAUNCH, PAD, 11)
+_FIELD = np.stack(
+    np.meshgrid(np.linspace(-1.0, 6.0, 8), np.linspace(-2.0, 3.0, 6)), axis=-1
+).reshape(-1, 2)
+
+
+def probe(model: LinearGaussianModel) -> ProbeReport:
+    """Probe ``model`` over the nine compass moves from its prior.
+
+    Args:
+        model: a rocket's generative model.
+
+    Returns:
+        cpomdp's report on whether a move can change what the rocket learns.
+    """
+    return _probe_model(model, model.prior, compass_moves())
+
+
+def epistemic_spread(report: ProbeReport) -> float:
+    """How far apart the best and worst move are in what they teach, in nats."""
+    lowest, highest = report.epistemic_range
+    return highest - lowest
+
+
+def show_probe(report: ProbeReport) -> None:
+    """Print the three numbers task 1 asks about, on one line."""
+    print(
+        f"epistemic_varies = {report.epistemic_varies}   "
+        f"flattens = {report.flattens}   "
+        f"spread = {epistemic_spread(report):.6f} nats"
+    )
+
+
+def check_agent(agent: Agent) -> None:
+    """Print the agent as built, then fly it in its own world and say where it lands.
+
+    The flight uses the agent up. Build it again to fly it again.
+
+    Args:
+        agent: a rocket that has not flown.
+    """
+    print(agent)
+    flight = fly(agent)
+    x, y = flight.final
+    print(
+        f"after {len(flight.path) - 1} moves it is at ({x:.2f}, {y:.2f}), "
+        f"{flight.distance_to(PAD):.2f} from the pad"
+    )
+
+
+@dataclass(frozen=True)
+class SensorGrade:
+    """What a sensor does for a rocket.
+
+    Attributes:
+        grade: ``"broken"``, ``"flat"``, ``"straight"``, ``"bends"`` or
+            ``"too curious"``.
+        message: what to tell the person who wrote it.
+        bend: how far the rocket strayed from the straight line, or ``None`` if it
+            did not fly.
+        distance: how far from the pad it ended, or ``None`` if it did not fly.
+        matches: the reference sensor it reproduces, or ``None``.
+    """
+
+    grade: str
+    message: str
+    bend: float | None = None
+    distance: float | None = None
+    matches: str | None = None
+
+
+def _untraceable(sensor: NoiseFn, params: SensorParams) -> str | None:
+    """Why JAX cannot trace ``sensor``, or ``None`` if it can.
+
+    The search differentiates and compiles through the sensor, so Python's ``min``,
+    the ``math`` module and NumPy fail there even where a plain call works.
+    """
+    try:
+        jax.jit(sensor)(jnp.asarray(_ROUTE[0]), params)
+    except Exception as error:  # the learner's code, any failure is the finding
+        return (
+            f"JAX cannot trace it ({type(error).__name__}). Use jnp functions such as "
+            "jnp.minimum, jnp.exp and jnp.sum in place of min, math or np."
+        )
+    return None
+
+
+def _broken(sensor: NoiseFn, params: SensorParams) -> str | None:
+    """Why ``sensor`` is not a noise covariance the search can use, or ``None``."""
+    for point in np.vstack([_ROUTE, _FIELD]):
+        where = f"({point[0]:.1f}, {point[1]:.1f})"
+        try:
+            noise = np.asarray(sensor(point, params), dtype=float)
+        except Exception as error:  # the learner's code, any failure is the finding
+            return f"calling it at {where} raised {type(error).__name__}: {error}"
+        if noise.shape != (2, 2):
+            return (
+                f"it returned shape {noise.shape} at {where}. It needs a 2x2 matrix, "
+                "jnp.eye(2) * (...)"
+            )
+        if not np.all(np.isfinite(noise)):
+            return f"it returned a non-finite value at {where}."
+        if not np.allclose(noise, noise.T) or np.linalg.eigvalsh(noise).min() <= 0:
+            return f"it is not a positive noise at {where}: {noise.tolist()}"
+    return _untraceable(sensor, params)
+
+
+def _matching_reference(sensor: NoiseFn, params: SensorParams) -> str | None:
+    """The name of the reference sensor ``sensor`` reproduces, if any."""
+    for name, reference in (
+        ("beacon_sensor", beacon_sensor),
+        ("reference_sensor", reference_sensor),
+    ):
+        if all(
+            np.allclose(sensor(p, params), reference(p, params), atol=1e-12)
+            for p in np.vstack([_ROUTE, _FIELD])
+        ):
+            return name
+    return None
+
+
+def grade_sensor(sensor: NoiseFn, params: SensorParams) -> SensorGrade:
+    """Grade a noise function by flying a curious rocket with it.
+
+    The rocket models ``sensor`` and flies in it from the default seed, as the
+    attempt cell's rocket does. A probe of what each move would teach is not the
+    grade, because it does not predict the path. The reading noise also enters what
+    each move is expected to read against "on the pad", and at these settings that
+    term does most of the steering.
+
+    Args:
+        sensor: a noise function ``R(x, params)``.
+        params: its parameters.
+
+    Returns:
+        A [`SensorGrade`][cpomdp_workshop.SensorGrade].
+    """
+    problem = _broken(sensor, params)
+    if problem is not None:
+        return SensorGrade("broken", f"Not a sensor yet: {problem}")
+    noises = [np.asarray(sensor(p, params)) for p in np.vstack([_ROUTE, _FIELD])]
+    if all(np.allclose(noise, noises[0], atol=1e-12) for noise in noises):
+        return SensorGrade(
+            "flat",
+            "R(x) is the same everywhere. That is the 1960 rocket's sensor, and no "
+            "move can teach it more than another.",
+        )
+    model = model_with(sensor, params)
+    planner = compass_planner(model)
+    flight = fly(Agent(model, pad_preference(), selector=planner))
+    bend = float(np.abs(flight.path[:, 1]).max())
+    distance = flight.distance_to(PAD)
+    matches = _matching_reference(sensor, params)
+    if bend < STRAIGHT_BEND:
+        message = (
+            "R(x) varies across the map. The rocket still flies straight, so the "
+            "pad's pull wins at this k. Raise it."
+        )
+        return SensorGrade("straight", message, bend, distance, matches)
+    path = f"The rocket bends {bend:.2f} off the straight line"
+    if distance > LANDED:
+        message = (
+            f"{path} and stops {distance:.2f} short of the pad. At this k, reading "
+            "near the beacon is worth more than reaching the pad."
+        )
+        return SensorGrade("too curious", message, bend, distance, matches)
+    message = f"{path} and lands {distance:.2f} from the pad."
+    if matches is not None:
+        message += f" It matches the workshop's {matches}."
+    return SensorGrade("bends", message, bend, distance, matches)
+
+
+def check_my_sensor(sensor: NoiseFn, params: SensorParams) -> None:
+    """Print what ``sensor`` does for a rocket, graded.
+
+    Args:
+        sensor: a noise function ``R(x, params)``.
+        params: its parameters, from ``default_params(k)``.
+    """
+    print(grade_sensor(sensor, params).message)
+
+
+BeliefTrack = tuple[np.ndarray, np.ndarray]
+
+
+def _belief_track(backend: InferenceBackend) -> BeliefTrack:
+    """The belief means and covariances of the 1960 flight run on ``backend``."""
+    flight = fly(Agent.from_backend(backend, pad_goal()))
+    return flight.means, flight.covs
+
+
+def _largest_difference(track: BeliefTrack, reference: BeliefTrack) -> float:
+    """The largest absolute difference in any belief mean or covariance entry."""
+    return float(
+        max(
+            np.abs(track[0] - reference[0]).max(), np.abs(track[1] - reference[1]).max()
+        )
+    )
+
+
+def backend_differences(
+    reference: InferenceBackend, *others: InferenceBackend
+) -> list[float]:
+    """How far each backend's beliefs stray from the reference's over one flight.
+
+    Every backend flies the 1960 rocket's goal from the same seed, in the world its
+    model describes.
+
+    Args:
+        reference: the backend the others are measured against.
+        others: the backends to measure.
+
+    Returns:
+        For each of ``others``, the largest absolute difference from the reference in
+        any belief mean or covariance entry, over every step.
+    """
+    reference_track = _belief_track(reference)
+    return [
+        _largest_difference(_belief_track(backend), reference_track)
+        for backend in others
+    ]
+
+
+def check_backends(reference: InferenceBackend, *others: InferenceBackend) -> None:
+    """Fly the 1960 rocket on each backend and print where their beliefs end up.
+
+    Args:
+        reference: the backend the others are measured against.
+        others: the backends to measure.
+    """
+    backends = (reference, *others)
+    tracks = [_belief_track(backend) for backend in backends]
+    for backend, (means, covs) in zip(backends, tracks, strict=True):
+        x, y = means[-1]
+        print(
+            f"{type(backend).__name__:<22} mean ({x:.6f}, {y:.6f})   "
+            f"cov diagonal ({covs[-1, 0, 0]:.6f}, {covs[-1, 1, 1]:.6f})"
+        )
+    largest = max(_largest_difference(track, tracks[0]) for track in tracks[1:])
+    print(f"max |difference| = {largest:.1e}")