mujax 0.1.0__py3-none-any.whl

mujax/__init__.py

@@ -0,0 +1,55 @@
+"""Stochastic MuZero and Gumbel MuZero, built on JAX, mctx and Gymnasium.
+
+Module map:
+    config.py      MuZeroConfig
+    algorithm.py   Algorithm, SearchPolicy: the contract between the infrastructure and an algorithm
+    algorithms/    ALGORITHMS registry and the algorithms
+        muzero.py      shared math of the MuZero family: Support, network blocks, loss skeleton
+        smz.py / gmz.py  algorithms (config, networks, search, loss)
+    replay.py / reanalyze.py  flashbax replay and reanalyze iterator
+    actor.py / learner.py / loop.py  environment interaction, gradient step, env loop
+    observers.py / loggers.py / checkpoint.py  metrics, terminal/W&B writers, orbax checkpoint
+    experiment.py  ExperimentConfig, run_experiment, load_actor
+"""
+
+from mujax.actor import Actor
+from mujax.algorithm import Algorithm
+from mujax.algorithms import ALGORITHMS, GMZ, SMZ, GMZConfig, SMZConfig, algorithm_for
+from mujax.checkpoint import CheckpointingConfig
+from mujax.config import MuZeroConfig
+from mujax.experiment import EvaluationConfig, ExperimentConfig, load_actor, run_experiment
+from mujax.learner import Learner
+from mujax.loggers import ConsoleWriter, Logger, WandbWriter, Writer, create_writer
+from mujax.loop import EnvironmentLoop
+from mujax.observers import EnvLoopObserver, InfoKeysObserver
+from mujax.types import EnvironmentSpec, make_environment_spec
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "ALGORITHMS",
+    "GMZ",
+    "SMZ",
+    "Actor",
+    "Algorithm",
+    "CheckpointingConfig",
+    "ConsoleWriter",
+    "EnvLoopObserver",
+    "EnvironmentLoop",
+    "EnvironmentSpec",
+    "EvaluationConfig",
+    "ExperimentConfig",
+    "GMZConfig",
+    "InfoKeysObserver",
+    "Learner",
+    "Logger",
+    "MuZeroConfig",
+    "SMZConfig",
+    "WandbWriter",
+    "Writer",
+    "algorithm_for",
+    "create_writer",
+    "load_actor",
+    "make_environment_spec",
+    "run_experiment",
+]

mujax/actor.py

@@ -0,0 +1,131 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Any
+
+import jax
+import numpy as np
+
+from mujax.algorithm import SearchPolicy
+from mujax.replay import Adder
+from mujax.types import SearchOutput, TimeStep, Transition
+
+
+def invalid_actions_from_info(info: dict | None, num_envs: int, num_actions: int) -> np.ndarray:
+    """If `info["invalid_actions"]` (0/1 mask, 1 = invalid) is absent, all actions are treated as valid."""
+    raw = info.get("invalid_actions") if info else None
+    if raw is None:
+        return np.zeros((num_envs, num_actions), dtype=np.float32)
+    return np.asarray(raw, dtype=np.float32).reshape(num_envs, num_actions)
+
+
+class Actor:
+    """The env-facing side: jitted search policy + learner params + writing to replay.
+
+    Observations are (num_envs, obs_dim), actions are (num_envs,).
+
+    Args:
+        get_params: () -> (params, learner_steps); usually `Learner.get_params`.
+        replay: if given (`Buffer.adder(i)`), every step is written there (training actor).
+        value_fn: (params, obs) -> search-free value (N,); for episodes cut by the time limit,
+            to bootstrap with V(final_obs). If None, a truncated episode is treated as terminated.
+        update_period: how many `select_action` calls between params fetches.
+        per_episode_update: if True, params are fetched only at episode start (evaluator: the score
+            matches the params that played that episode).
+    """
+
+    def __init__(
+        self,
+        policy: SearchPolicy,
+        random_key: jax.Array,
+        get_params: Callable[[], tuple[Any, int]],
+        *,
+        num_actions: int,
+        replay: Adder | None = None,
+        value_fn: Callable[[Any, Any], Any] | None = None,
+        device: jax.Device | None = None,
+        update_period: int = 1,
+        per_episode_update: bool = False,
+    ) -> None:
+        self._search = jax.jit(policy)
+        self._key = random_key
+        self._get_params = get_params
+        self._num_actions = num_actions
+        self._replay = replay
+        self._value = jax.jit(value_fn) if value_fn is not None else None
+        self._device = device
+        self._update_period = max(1, update_period)
+        self._per_episode_update = per_episode_update
+        self._num_calls = 0
+        self._episode_start = True
+        self._last_search: SearchOutput | None = None
+        self._invalid_actions: np.ndarray | None = None
+        self._pull_params()
+
+    def _pull_params(self) -> None:
+        params, learner_steps = self._get_params()
+        self._params = (jax.device_put(params, self._device), learner_steps)  # single assignment
+
+    def get_params(self) -> tuple[Any, int]:
+        """The (params, learner_steps) currently playing."""
+        return self._params
+
+    def observe_first(self, observation: np.ndarray, info: dict) -> None:
+        self._invalid_actions = invalid_actions_from_info(info, len(observation), self._num_actions)
+        self._episode_start = True
+
+    def select_action(self, observation: np.ndarray) -> np.ndarray:
+        self._num_calls += 1
+        if (self._per_episode_update and self._episode_start) or (not self._per_episode_update and self._num_calls % self._update_period == 0):
+            self._pull_params()
+        self._episode_start = False
+
+        params, learner_steps = self._params
+        self._key, search_key = jax.random.split(self._key)
+        search = self._search(params, np.asarray(observation), search_key, self._invalid_actions, learner_steps)
+        self._last_search = SearchOutput(
+            action=np.asarray(search.action, dtype=np.int32),
+            policy_probs=np.asarray(search.policy_probs, dtype=np.float32),
+            value=np.asarray(search.value, dtype=np.float32),
+        )
+        return self._last_search.action
+
+    def _final_values(self, step: TimeStep, truncated: np.ndarray) -> np.ndarray:
+        """V(final_obs) for truncated envs; under SAME_STEP the final observation is in `info["final_obs"]`."""
+        values = np.zeros(len(truncated), dtype=np.float32)
+        if not truncated.any():
+            return values
+        # To call the jit with a single shape, the whole batch is filled; non-truncated rows are discarded.
+        obs = np.zeros_like(np.asarray(step.observation), dtype=np.float32)
+        for index in np.flatnonzero(truncated):
+            obs[index] = step.info["final_obs"][index]
+        values[truncated] = np.asarray(self._value(self._params[0], obs))[truncated]
+        return values
+
+    def last_search(self) -> SearchOutput | None:
+        """Search output of the last `select_action`; EnvironmentLoop puts it into the TimeStep."""
+        return self._last_search
+
+    def observe(self, step: TimeStep) -> None:
+        terminated = np.asarray(step.terminated, dtype=bool)
+        truncated = np.asarray(step.truncated, dtype=bool) & ~terminated
+        done = terminated | truncated
+        if self._replay is not None:
+            if self._value is None:  # V(final_obs) cannot be computed: a truncated episode is treated as terminated
+                terminated, truncated = done, np.zeros_like(truncated)
+            self._replay.add(
+                Transition(
+                    observation=np.asarray(step.observation, dtype=np.float32),
+                    action=self._last_search.action,
+                    reward=np.asarray(step.reward, dtype=np.float32),
+                    discount=np.where(terminated, 0.0, 1.0).astype(np.float32),
+                    policy_probs=self._last_search.policy_probs,
+                    value=self._last_search.value,
+                    invalid_actions=self._invalid_actions,
+                    truncated=truncated.astype(np.float32),
+                    bootstrap_value=self._final_values(step, truncated),
+                )
+            )
+        self._invalid_actions = invalid_actions_from_info(step.info, len(done), self._num_actions)
+        if done.any():
+            self._episode_start = True

mujax/algorithm.py

@@ -0,0 +1,53 @@
+"""The contract between the shared infrastructure and an algorithm in `mujax.algorithms`."""
+
+from __future__ import annotations
+
+import dataclasses
+from collections.abc import Callable
+from typing import Any
+
+import jax
+import jax.numpy as jnp
+from flax import nnx
+
+from mujax.config import MuZeroConfig
+from mujax.types import EnvironmentSpec, SearchOutput
+
+# Batched search policy: (params, obs, key, invalid_actions, learner_steps) -> SearchOutput.
+# All inputs and outputs have shape (num_envs, ...).
+SearchPolicy = Callable[[Any, Any, jax.Array, Any, Any], SearchOutput]
+
+# Search-free root value: (params, obs) -> (N,); used for the bootstrap tail in reanalyze.
+ValueFn = Callable[[Any, Any], jnp.ndarray]
+
+# Algorithm loss: (graphdef, config, params, batch) -> (loss, metrics).
+LossFn = Callable[[Any, MuZeroConfig, Any, Any], tuple[jnp.ndarray, dict[str, jnp.ndarray]]]
+
+
+@dataclasses.dataclass(frozen=True)
+class Algorithm:
+    """Algorithm-specific parts; everything else (replay, reanalyze, actor, learner, loop) is shared.
+
+    The model is a single `nnx.Module`; `init` splits it into a graphdef (static structure) and params
+    (`nnx.State`, a pytree). Params travel separately (learner, actor threads, mctx, checkpoints) and the
+    algorithm functions rebuild the model with `nnx.merge(graphdef, params)`.
+
+    Attributes:
+        name: Short name written to checkpoint metadata ("gmz", "smz").
+        config_cls: Config class of this algorithm; `algorithm_for` uses it to pick the algorithm from a config.
+        make_model: (spec, config, rngs) -> model (nnx.Module).
+        make_policy: (graphdef, spec, config, evaluation) -> SearchPolicy.
+        loss: See `LossFn`.
+        make_value_fn: (graphdef, config) -> ValueFn.
+    """
+
+    name: str
+    config_cls: type[MuZeroConfig]
+    make_model: Callable[[EnvironmentSpec, MuZeroConfig, nnx.Rngs], nnx.Module]
+    make_policy: Callable[[nnx.GraphDef, EnvironmentSpec, MuZeroConfig, bool], SearchPolicy]
+    loss: LossFn
+    make_value_fn: Callable[[nnx.GraphDef, MuZeroConfig], ValueFn]
+
+    def init(self, spec: EnvironmentSpec, config: MuZeroConfig, key: jax.Array) -> tuple[nnx.GraphDef, nnx.State]:
+        """(graphdef, params) of a freshly initialized model."""
+        return nnx.split(self.make_model(spec, config, nnx.Rngs(key)))

mujax/algorithms/__init__.py

@@ -0,0 +1,21 @@
+"""Algorithms; each module defines an `Algorithm` (see `mujax.algorithm`). New algorithms are registered in `ALGORITHMS`."""
+
+from __future__ import annotations
+
+from mujax.algorithm import Algorithm
+from mujax.algorithms.gmz import GMZ, GMZConfig
+from mujax.algorithms.smz import SMZ, SMZConfig
+from mujax.config import MuZeroConfig
+
+ALGORITHMS: dict[str, Algorithm] = {SMZ.name: SMZ, GMZ.name: GMZ}
+
+
+def algorithm_for(config: MuZeroConfig) -> Algorithm:
+    """The algorithm whose config class `config` is an instance of."""
+    for algorithm in ALGORITHMS.values():
+        if isinstance(config, algorithm.config_cls):
+            return algorithm
+    raise ValueError(f"{type(config).__name__} icin algoritma bulunamadi")
+
+
+__all__ = ["ALGORITHMS", "GMZ", "SMZ", "GMZConfig", "SMZConfig", "algorithm_for"]

mujax/algorithms/gmz.py

@@ -0,0 +1,116 @@
+"""Gumbel MuZero: config, networks, search and loss."""
+
+from __future__ import annotations
+
+import dataclasses
+import functools
+from collections.abc import Sequence
+
+import jax.numpy as jnp
+import mctx
+from flax import nnx
+
+from mujax.algorithm import Algorithm, SearchPolicy
+from mujax.algorithms import muzero
+from mujax.algorithms.muzero import Embedding, Head, Prediction, Support, SymlogInput, Targets
+from mujax.config import MuZeroConfig
+from mujax.types import EnvironmentSpec, SearchOutput
+
+
+@dataclasses.dataclass
+class GMZConfig(MuZeroConfig):
+    """Gumbel MuZero hyperparameters."""
+
+    num_simulations: int = 16
+    max_depth: int | None = None
+    max_num_considered_actions: int = 16  # clipped to the number of actions
+    # Gumbel noise is gmz's only source of exploration (smz uses temperature + Dirichlet).
+    gumbel_scale: float = 1.0
+    value_scale: float = 0.1
+    maxvisit_init: float = 50.0
+
+    dynamics_layer_sizes: tuple[int, ...] = (256, 256, 256)
+
+
+class Dynamics(nnx.Module):
+    def __init__(self, embedding_dim: int, layer_sizes: Sequence[int], num_actions: int, num_bins: int, *, rngs: nnx.Rngs):
+        self.num_actions = num_actions
+        self.embedding = Embedding(embedding_dim + num_actions, layer_sizes, embedding_dim, rngs=rngs)
+        self.reward = Head(embedding_dim, layer_sizes, num_bins, rngs=rngs)
+
+    def __call__(self, latent, action) -> tuple[jnp.ndarray, jnp.ndarray]:
+        """(next_latent, reward_logits)."""
+        next_latent = self.embedding(muzero.one_hot_concat(latent, action, self.num_actions))
+        return next_latent, self.reward(next_latent)
+
+
+class GMZModel(nnx.Module):
+    """The deterministic model: representation, prediction and dynamics."""
+
+    def __init__(self, spec: EnvironmentSpec, config: GMZConfig, rngs: nnx.Rngs):
+        self.representation = SymlogInput(Embedding(spec.obs_dim, config.representation_layer_sizes, config.embedding_dim, rngs=rngs))
+        self.prediction = Prediction(
+            config.embedding_dim, config.prediction_layer_sizes, spec.num_actions, config.num_bins, unimix=config.policy_unimix, rngs=rngs
+        )
+        self.dynamics = Dynamics(config.embedding_dim, config.dynamics_layer_sizes, spec.num_actions, config.num_bins, rngs=rngs)
+
+
+def recurrent_fn(graphdef: nnx.GraphDef, support: Support, config: GMZConfig):
+    """mctx recurrent_fn of the deterministic model."""
+
+    def fn(params, key, action, latent):
+        model = nnx.merge(graphdef, params)
+        next_latent, reward_logits = model.dynamics(latent, action)
+        value_logits, policy_logits = model.prediction(next_latent)
+        reward = support.logits_to_scalar(reward_logits)
+        output = mctx.RecurrentFnOutput(
+            reward=reward,
+            discount=jnp.full_like(reward, config.discount),
+            prior_logits=policy_logits,
+            value=support.logits_to_scalar(value_logits),
+        )
+        return output, next_latent
+
+    return fn
+
+
+def make_policy(graphdef: nnx.GraphDef, spec: EnvironmentSpec, config: GMZConfig, evaluation: bool) -> SearchPolicy:
+    """Batched Gumbel MuZero search. `evaluation=True` turns off the Gumbel noise."""
+    support = Support.from_config(config)
+    max_num_considered_actions = min(int(config.max_num_considered_actions), spec.num_actions)
+    qtransform = functools.partial(mctx.qtransform_completed_by_mix_value, value_scale=config.value_scale, maxvisit_init=config.maxvisit_init)
+
+    def policy(params, obs, key, invalid_actions, learner_steps) -> SearchOutput:
+        del learner_steps
+        out = mctx.gumbel_muzero_policy(
+            params,
+            key,
+            muzero.root_output(nnx.merge(graphdef, params), support, obs),
+            recurrent_fn(graphdef, support, config),
+            num_simulations=config.num_simulations,
+            invalid_actions=invalid_actions,
+            max_depth=config.max_depth,
+            qtransform=qtransform,
+            max_num_considered_actions=max_num_considered_actions,
+            gumbel_scale=0.0 if evaluation else config.gumbel_scale,
+        )
+        return SearchOutput(out.action, out.action_weights, out.search_tree.summary().value)
+
+    return policy
+
+
+def loss(graphdef: nnx.GraphDef, config: GMZConfig, params: nnx.State, batch) -> tuple[jnp.ndarray, dict[str, jnp.ndarray]]:
+    def unroll_step(model: GMZModel, latent, now: Targets, next_: Targets):
+        next_latent, reward_logits = model.dynamics(latent, now.action)
+        value_logits, policy_logits = model.prediction(next_latent)
+        losses = {
+            "reward": muzero.masked_mean(muzero.cross_entropy(reward_logits, now.reward_probs), now.in_episode),
+            "value": muzero.masked_mean(muzero.cross_entropy(value_logits, next_.value_probs), next_.in_episode),
+            "policy": muzero.masked_mean(muzero.cross_entropy(policy_logits, next_.policy_probs), next_.in_episode),
+        }
+        return next_latent, losses, {}
+
+    return muzero.loss(graphdef, config, params, batch, unroll_step=unroll_step, loss_weights={"value": config.value_loss_weight})
+
+
+GMZ = Algorithm("gmz", GMZConfig, GMZModel, make_policy, loss, muzero.make_value_fn)