Skip to content

Process noise

State-dependent dynamics noise Q(x): the world can diffuse more in some states than others. CallableProcessNoise carries the Q(x) function; both it and a plain constant satisfy the DynamicsNoise protocol.

DynamicsNoise

Bases: Protocol

How much the dynamics diffuse per step — the process noise covariance Q(x).

The EFE predict step asks for Q at a state. is_fixed flags the constant case (a plain matrix); a state-dependent Q(x) re-introduces epistemic value through the internal route (RFC-001 Section8).

noise_at 

noise_at(x: ArrayLike) -> Float64[Array, 'n n']

The process-noise covariance Q(x) at state x.

Source code in src/cpomdp/dynamics.py 
32
33
34
def noise_at(self, x: ArrayLike) -> Float64[Array, "n n"]:
    """The process-noise covariance ``Q(x)`` at state ``x``."""
    ...

CallableProcessNoise dataclass 

CallableProcessNoise(q_fn, q_params)

State-dependent process noise Q(x) = q_fn(x, q_params).

q_fn holds the functional form (static aux — a callable can't be a traced leaf); q_params holds the values it is a function of (a pytree leaf, so EFE is grad-able w.r.t. them — process-noise learning). The rule for what goes where: a number lives in q_params if you'd ever want to fit it; it's baked into q_fn only if it's structural and never learned. Pass a module-level q_fn (a closure/lambda hashes by identity and defeats jit caching).

Source code in src/cpomdp/dynamics.py 
54
55
56
57
def __init__(self, q_fn, q_params) -> None:
    object.__setattr__(self, "q_fn", q_fn)
    object.__setattr__(self, "q_params", q_params)
    self._validate()

noise_at 

noise_at(x: ArrayLike) -> Float64[Array, 'n n']

Q(x) — the process-noise covariance at state x.

Source code in src/cpomdp/dynamics.py 
59
60
61
def noise_at(self, x: ArrayLike) -> Float64[Array, "n n"]:
    """Q(x) — the process-noise covariance at state x."""
    return self.q_fn(x, self.q_params)

tree_flatten 

tree_flatten() -> tuple[tuple[PyTree], Callable]

Children (traced): (q_params,); aux (static): q_fn.

Source code in src/cpomdp/dynamics.py 
63
64
65
def tree_flatten(self) -> tuple[tuple[PyTree], Callable]:
    """Children (traced): ``(q_params,)``; aux (static): ``q_fn``."""
    return (self.q_params,), self.q_fn

tree_unflatten classmethod 

tree_unflatten(
    aux_data: Callable, children: tuple
) -> CallableProcessNoise

Rebuild without re-validating — leaves may be tracers.

Source code in src/cpomdp/dynamics.py 
67
68
69
70
71
72
73
74
75
76
@classmethod
def tree_unflatten(
    cls, aux_data: Callable, children: tuple
) -> "CallableProcessNoise":
    """Rebuild without re-validating — leaves may be tracers."""
    (q_params,) = children
    obj = object.__new__(cls)
    object.__setattr__(obj, "q_fn", aux_data)
    object.__setattr__(obj, "q_params", q_params)
    return obj