MASA-Safe-RL 0.3.2__tar.gz → 0.3.4__tar.gz

{masa_safe_rl-0.3.2 → masa_safe_rl-0.3.4/MASA_Safe_RL.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: MASA-Safe-RL
-Version: 0.3.2
+Version: 0.3.4
 Summary: A safe reinforcement library for providing a common interface for different constraints and environments
 Project-URL: Homepage, https://github.com/sacktock/MASA-Safe-RL
 Project-URL: Repository, https://github.com/sacktock/MASA-Safe-RL

{masa_safe_rl-0.3.2 → masa_safe_rl-0.3.4}/MASA_Safe_RL.egg-info/SOURCES.txt

@@ -188,6 +188,7 @@ masa/plotting/plot.py
 masa/plotting/processing.py
 masa/plotting/render.py
 masa/plotting/run.py
+masa/plotting/source_registry.py
 masa/plotting/sources.py
 masa/plotting/specs/__init__.py
 masa/plotting/specs/base.py

{masa_safe_rl-0.3.2/MASA_Safe_RL.egg-info → masa_safe_rl-0.3.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: MASA-Safe-RL
-Version: 0.3.2
+Version: 0.3.4
 Summary: A safe reinforcement library for providing a common interface for different constraints and environments
 Project-URL: Homepage, https://github.com/sacktock/MASA-Safe-RL
 Project-URL: Repository, https://github.com/sacktock/MASA-Safe-RL

{masa_safe_rl-0.3.2 → masa_safe_rl-0.3.4}/masa/algorithms/on_policy/trpo.py

@@ -30,10 +30,10 @@ class TRPO(OnPolicyAlgorithm):
         vf_coef: float = 1.0,
         max_grad_norm: float = 0.5,
         target_kl: float = 0.01,
-        cg_iters: int = 10,
+        cg_iters: int = 15,
         cg_damping: float = 0.1,
-        fvp_sample_freq: int = 4,
-        line_search_steps: int = 10,
+        fvp_sample_freq: int = 1,
+        line_search_steps: int = 15,
         line_search_decay: float = 0.8,
         policy_class: type[BaseJaxPolicy] = PPOPolicy,
         policy_kwargs: Optional[dict[str, Any]] = None,

{masa_safe_rl-0.3.2 → masa_safe_rl-0.3.4}/masa/algorithms/tabular/recreg.py

@@ -512,8 +512,13 @@ class RECREG(QL):
 
         if override:
             # Use boltzmann temperature=0.01 or epsilon=0.001 for backup policy
-            safe_act = self._act(key3, obs, self.B, deterministic, 0.01, 0.001)
-
+            if self.exploration == "boltzmann":
+                safe_act = self._act(key3, obs, self.B, deterministic, 0.01, 0.001)
+            elif self.exploration == "epsilon_greedy":
+                safe_act = self._act(key3, obs, self.B, True, 0.01, 0.001)
+            else:
+                raise NotImplementedError(f"Unexpected exploration: {self.exploration}")
+    
         return self.prepare_act(safe_act), self.prepare_act(act), override
 
     def act(self, key, obs, deterministic=False):

{masa_safe_rl-0.3.2 → masa_safe_rl-0.3.4}/masa/configs/algorithms/CPO.yaml

@@ -50,11 +50,6 @@ ColourBombGridWorldV2:
   cost_limit: 0.025
   cost_gamma: 0.95
 
-ColourBombGridWorldV2_Property1:
-  gamma: 0.95
-  cost_limit: 0.025
-  cost_gamma: 0.95
-
 ColourBombGridWorldV2_Property2:
   gamma: 0.95
   cost_limit: 0.025

{masa_safe_rl-0.3.2 → masa_safe_rl-0.3.4}/masa/configs/algorithms/PPOLag.yaml

@@ -55,12 +55,6 @@ ColourBombGridWorldV2:
   cost_gamma: 0.95
   lagrangian_multiplier_init: 10.0
 
-ColourBombGridWorldV2_Property1:
-  gamma: 0.95
-  cost_limit: 0.025
-  cost_gamma: 0.95
-  lagrangian_multiplier_init: 10.0
-
 ColourBombGridWorldV2_Property2:
   gamma: 0.95
   cost_limit: 0.025

{masa_safe_rl-0.3.2 → masa_safe_rl-0.3.4}/masa/configs/algorithms/ProbShieldPPO.yaml

@@ -10,7 +10,7 @@ defaults:
   ent_coef: 0.0
   vf_coef: 1.0
   max_grad_norm: 0.5
-  policy_kwargs: {net_arch: [64, 64], log_std_init: -1.0}
+  policy_kwargs: {net_arch: [64, 64], log_std_init: -1.0, conditional_beta_network: True}
   init_safety_bound: 0.01
   theta: 1e-15
   max_vi_steps: 10000
@@ -39,11 +39,6 @@ ColourBombGridWorldV2:
   gamma: 0.95
   init_safety_bound: 0.025
 
-ColourBombGridWorldV2_Property1:
-  gamma: 0.95
-  init_safety_bound: 0.025
-  policy_kwargs: {log_std_init: -2.0}
-
 ColourBombGridWorldV2_Property2:
   gamma: 0.95
   init_safety_bound: 0.025

{masa_safe_rl-0.3.2 → masa_safe_rl-0.3.4}/masa/configs/algorithms/RECREG_EXACT.yaml

@@ -114,7 +114,7 @@ MiniPacman:
   safe_gamma: 0.95
   mode: exact
   model_checking: exact
-  horizon: 9
+  horizon: 20
   safety_prob: 0.01
   cost_coef: 10.0
   

{masa_safe_rl-0.3.2 → masa_safe_rl-0.3.4}/masa/configs/algorithms/RECREG_MODEL_BASED.yaml

@@ -114,7 +114,7 @@ MiniPacman:
   safe_gamma: 0.95
   mode: model_based
   model_checking: exact
-  horizon: 9
+  horizon: 20
   safety_prob: 0.01
   cost_coef: 10.0
   

{masa_safe_rl-0.3.2 → masa_safe_rl-0.3.4}/masa/configs/algorithms/recreg.yaml

@@ -114,7 +114,7 @@ MiniPacman:
   safe_gamma: 0.95
   model_free_alpha: 0.1
   mode: model_free
-  horizon: 9
+  horizon: 20
   safety_prob: 0.01
   cost_coef: 10.0
   

{masa_safe_rl-0.3.2 → masa_safe_rl-0.3.4}/masa/plotting/__main__.py

@@ -1,17 +1,7 @@
-"""CLI entry point.
+"""Benchmark pipeline CLI: load runs, aggregate quantiles, render PlotSpecs.
 
-Example (generic researcher, ships with the ``performance`` spec):
-
-    python -m masa.plotting \\
-        --config masa/plotting/configs/example.yaml \\
-        --stages all
-
-Example (load extra specs from another module before resolving):
-
-    python -m masa.plotting \\
-        --config <your-yaml> \\
-        --specs-from <your.spec.module> \\
-        --stages process,render
+    python -m masa.plotting --config configs/example_wandb.yaml --stages all
+    python -m masa.plotting --config <yaml> --specs-from <module> --specs my_plot
 """
 
 from __future__ import annotations
@@ -55,10 +45,10 @@ def _preload_spec_modules(modules: Sequence[str]) -> None:
 def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace:
     p = argparse.ArgumentParser(
         prog="python -m masa.plotting",
-        description="Download W&B runs, aggregate to quantiles, render figures from PlotSpecs.",
+        description="Load runs from W&B or TensorBoard, aggregate to quantiles, render figures from PlotSpecs.",
     )
     p.add_argument("--config", required=True, type=Path,
-                   help="Path to a YAML config (see configs/example.yaml).")
+                   help="Path to a YAML config (see configs/example_wandb.yaml or example_tensorboard.yaml).")
     p.add_argument("--stages", default="all",
                    help="Comma separated subset of {download,process,render}, or 'all'.")
     p.add_argument("--specs", default=None,

{masa_safe_rl-0.3.2 → masa_safe_rl-0.3.4}/masa/plotting/api.py

@@ -2,11 +2,13 @@
 
 from __future__ import annotations
 
+import os
 from typing import Dict, List, Optional, Tuple
 
 import matplotlib.pyplot as plt
 
-from masa.plotting.plot import plot_metrics
+from masa.plotting.plot import plot_metrics, plot_quantile_metrics
+from masa.plotting.processing import build_long_form_from_tb, build_quantile_frame
 from masa.plotting.sources import TensorBoardSource, WandBSource
 
 
@@ -17,7 +19,8 @@ def plot_run(
     wandb_project: Optional[str] = None,
     wandb_entity: Optional[str] = None,
     wandb_filters: Optional[dict] = None,
-    run_filter: Optional[str] = None,
+    run_filter: Optional[str | List[str]] = None,
+    group_seeds: bool = False,
     # What to plot
     metrics: Optional[List[str]] = None,
     # Plotting options
@@ -66,9 +69,29 @@ def plot_run(
     if not tensorboard_logdir and not wandb_project:
         raise ValueError("Provide either tensorboard_logdir or wandb_project.")
 
+    if group_seeds and not tensorboard_logdir:
+        raise ValueError("group_seeds is only supported for TensorBoard sources.")
+
     if tensorboard_logdir:
         source = TensorBoardSource(tensorboard_logdir)
-        df = source.load(metrics=metrics, run_filter=run_filter)
+        raw_df = source.load(metrics=metrics, run_filter=run_filter)
+        if group_seeds:
+            env = os.path.basename(tensorboard_logdir.rstrip("/\\"))
+            long_df = build_long_form_from_tb(raw_df, env=env)
+            q_df = build_quantile_frame(long_df)
+            return plot_quantile_metrics(
+                q_df,
+                metrics=metrics,
+                smooth_weight=smooth_weight,
+                title=title,
+                output_path=output_path,
+                figsize=figsize,
+                show_legend=show_legend,
+                show_axes=show_axes,
+                dpi=dpi,
+                palette=palette,
+            )
+        df = raw_df
     else:
         source = WandBSource(wandb_project, entity=wandb_entity)
         df = source.load(metrics=metrics, filters=wandb_filters, run_filter=run_filter)

{masa_safe_rl-0.3.2 → masa_safe_rl-0.3.4}/masa/plotting/config.py

@@ -7,6 +7,10 @@ from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Optional
 
+# Default run-directory schema for TensorBoard sources: the leaf run dir is
+# named ``{variant}_{seed}`` (env comes from the parent directory).
+_TB_RUN_SCHEMA = r"^(?P<variant>.+)_(?P<seed>\d+)$"
+
 
 @dataclass(frozen=True)
 class VariantStyle:
@@ -28,17 +32,20 @@ class RunSchema:
 
 @dataclass(frozen=True)
 class Config:
-    wandb_entity: str
-    wandb_project: str
     run_schema: RunSchema
     variants: list[VariantStyle]
     cache_dir: Path
     output_dir: Path
+    source_kind: str = "wandb"  # "wandb" or "tensorboard"
+    wandb_entity: Optional[str] = None
+    wandb_project: Optional[str] = None
+    tensorboard_logdir: Optional[Path] = None
     metrics: Optional[list[str]] = None
     envs: Optional[list[str]] = None
     palette: str = "colorblind"
     force_download: bool = False
     seed_for_logged_quantiles: int = 1
+    max_step: Optional[int] = None
 
     def variant_by_name(self) -> dict[str, VariantStyle]:
         return {v.name: v for v in self.variants}
@@ -77,24 +84,68 @@ def load_config(path: str | Path) -> Config:
     ]
     variants = _fill_palette(variants, raw.get("palette", "colorblind"))
 
+    source_kind, wandb_entity, wandb_project, tb_logdir = _resolve_source(path, raw, _resolve)
+
     schema_raw = raw.get("run_schema") or {}
-    schema = RunSchema(pattern=schema_raw.get("pattern", RunSchema.pattern))
+    if "pattern" in schema_raw:
+        schema = RunSchema(pattern=schema_raw["pattern"])
+    elif source_kind == "tensorboard":
+        schema = RunSchema(pattern=_TB_RUN_SCHEMA)
+    else:
+        schema = RunSchema()
+
+    output_dir = _resolve(raw["output_dir"])
+    cache_raw = raw.get("cache_dir")
+    cache_dir = _resolve(cache_raw) if cache_raw else output_dir / "_cache"
 
     return Config(
-        wandb_entity=raw["wandb_entity"],
-        wandb_project=raw["wandb_project"],
         run_schema=schema,
         variants=variants,
-        cache_dir=_resolve(raw["cache_dir"]),
-        output_dir=_resolve(raw["output_dir"]),
+        cache_dir=cache_dir,
+        output_dir=output_dir,
+        source_kind=source_kind,
+        wandb_entity=wandb_entity,
+        wandb_project=wandb_project,
+        tensorboard_logdir=tb_logdir,
         metrics=raw.get("metrics"),
         envs=raw.get("envs"),
         palette=raw.get("palette", "colorblind"),
         force_download=raw.get("force_download", False),
         seed_for_logged_quantiles=raw.get("seed_for_logged_quantiles", 1),
+        max_step=raw.get("max_step"),
     )
 
 
+def _resolve_source(
+    path: Path, raw: dict, resolve
+) -> tuple[str, Optional[str], Optional[str], Optional[Path]]:
+    """Return ``(source_kind, wandb_entity, wandb_project, tensorboard_logdir)``.
+
+    Accepts an explicit ``source:`` block, or falls back to legacy top-level
+    ``wandb_entity`` / ``wandb_project`` keys (treated as a W&B source).
+    """
+    src = raw.get("source")
+    if src:
+        kind = src.get("kind", "wandb")
+        if kind == "wandb":
+            project = src.get("project")
+            if not project:
+                raise ValueError(f"{path}: source.kind=wandb requires source.project.")
+            return "wandb", src.get("entity"), project, None
+        if kind == "tensorboard":
+            logdir = src.get("logdir")
+            if not logdir:
+                raise ValueError(f"{path}: source.kind=tensorboard requires source.logdir.")
+            return "tensorboard", None, None, resolve(logdir)
+        raise ValueError(
+            f"{path}: unknown source.kind={kind!r}; expected 'wandb' or 'tensorboard'."
+        )
+    project = raw.get("wandb_project")
+    if not project:
+        raise ValueError(f"{path}: provide a 'source:' block, or top-level 'wandb_project'.")
+    return "wandb", raw.get("wandb_entity"), project, None
+
+
 def _fill_palette(variants: list[VariantStyle], palette_name: str) -> list[VariantStyle]:
     """Assign colours from a seaborn palette to any variant where colour is None."""
     missing_idx = [i for i, v in enumerate(variants) if v.colour is None]

{masa_safe_rl-0.3.2 → masa_safe_rl-0.3.4}/masa/plotting/pipeline.py

@@ -9,8 +9,7 @@ import pandas as pd
 
 from .config import Config
 from .io_utils import get_logger
-from .processing import build_frames
-from .sources import CachedWandbSource
+from .source_registry import get_source
 from .specs import PlotSpec, all_specs
 from .specs.base import RenderContext
 
@@ -21,19 +20,17 @@ class Pipeline:
     def __init__(self, config: Config):
         self.config = config
         self.log = get_logger()
-
-    # -- individual stages ------------------------------------------------
+        self.source = get_source(config)
 
     def download(self) -> list[Path]:
         try:
-            source = CachedWandbSource(self.config)
-            return source.download()
+            return self.source.download()
         except Exception as e:
             self.log.error(f"download stage crashed: {e}")
             return []
 
     def process(self) -> tuple[pd.DataFrame, pd.DataFrame]:
-        return build_frames(self.config)
+        return self.source.build_frames()
 
     def render(self,
                long_df: pd.DataFrame,
@@ -68,8 +65,6 @@ class Pipeline:
                     outputs.append(path)
         return outputs
 
-    # -- orchestration ----------------------------------------------------
-
     def run(self,
             specs: Optional[Sequence[PlotSpec]] = None,
             stages: Sequence[str] = VALID_STAGES,
@@ -104,8 +99,6 @@ class Pipeline:
 
         return []
 
-    # -- helpers ----------------------------------------------------------
-
     def _select_envs(self, discovered: list[str]) -> list[str]:
         """Restrict to ``config.envs`` if set; otherwise return all discovered envs."""
         if not self.config.envs:
@@ -128,23 +121,14 @@ class Pipeline:
     def _dry_run(self, specs: Optional[Sequence[PlotSpec]]) -> None:
         cfg = self.config
         self.log.info("=== DRY RUN ===")
-        self.log.info(f"W&B          : {cfg.wandb_entity}/{cfg.wandb_project}")
-        self.log.info(f"cache_dir    : {cfg.cache_dir}")
+        self.log.info(f"source       : {self.source.describe()}")
         self.log.info(f"output_dir   : {cfg.output_dir}")
         self.log.info(f"variants     : {[v.name for v in cfg.variants]}")
         self.log.info(f"palette      : {cfg.palette}")
         chosen = specs if specs is not None else all_specs()
         self.log.info(f"specs        : {[s.id for s in chosen]}")
-        envs: list[str] = []
-        if cfg.cache_dir.exists():
-            pat = cfg.run_schema.compile()
-            for p in sorted(cfg.cache_dir.glob("*.csv")):
-                m = pat.match(p.stem)
-                if m:
-                    envs.append(m.group("env"))
-            envs = sorted(set(envs))
-        envs = self._select_envs(envs)
-        self.log.info(f"envs (cache) : {envs or '(none discovered yet)'}")
+        envs = self._select_envs(self.source.discover_envs())
+        self.log.info(f"envs         : {envs or '(none discovered yet)'}")
         for spec in chosen:
             for env in envs or ["<env>"]:
                 target = cfg.output_dir / spec.out_subdir / f"{spec.id}_{env}.png"

{masa_safe_rl-0.3.2 → masa_safe_rl-0.3.4}/masa/plotting/plot.py

@@ -58,7 +58,6 @@ def _thousands_formatter() -> ticker.FuncFormatter:
     )
 
 
-# Core plotting
 def plot_metrics(
     df: pd.DataFrame,
     metrics: Optional[List[str]] = None,
@@ -102,6 +101,45 @@ def plot_metrics(
         )
 
 
+def plot_quantile_metrics(
+    q_df: pd.DataFrame,
+    metrics: Optional[List[str]] = None,
+    smooth_weight: float = 0.6,
+    title: Optional[str] = None,
+    output_path: Optional[str] = None,
+    figsize: Tuple[int, int] = (7, 5),
+    show_legend: bool = True,
+    show_axes: bool = True,
+    dpi: int = 300,
+    palette: Optional[Dict[str, str]] = None,
+) -> plt.Figure:
+    """Plot q50 ± q25/q75 bands from a quantile DataFrame.
+
+    Args:
+        q_df: DataFrame with columns ``[variant, step, metric, q25, q50, q75]``
+            as produced by :func:`~masa.plotting.processing.build_quantile_frame`.
+        metrics: Metric keys to plot.  ``None`` plots all found in *q_df*.
+        smooth_weight: EMA weight applied per ``(variant, metric)`` group.
+        title: Per-panel title override.  ``None`` uses the metric name.
+        output_path: Save path.  ``None`` skips saving.
+        figsize: ``(width, height)`` per subplot panel.
+        show_legend: Render a shared legend below the panels.
+        show_axes: Label axes.
+        dpi: Resolution for saved figures.
+        palette: Colour mapping ``{variant: colour}``.  ``None`` uses colorblind.
+
+    Returns:
+        The matplotlib :class:`~matplotlib.figure.Figure`.
+    """
+    with plt.rc_context(_RC_PARAMS):
+        return _plot_quantile_inner(
+            q_df, metrics=metrics, smooth_weight=smooth_weight,
+            title=title, output_path=output_path, figsize=figsize,
+            show_legend=show_legend, show_axes=show_axes, dpi=dpi,
+            palette=palette,
+        )
+
+
 def _plot_metrics_inner(
     df: pd.DataFrame,
     *,
@@ -185,7 +223,84 @@ def _plot_metrics_inner(
     return fig
 
 
-# Pretty-printing helpers
+def _plot_quantile_inner(
+    q_df: pd.DataFrame,
+    *,
+    metrics: Optional[List[str]],
+    smooth_weight: float,
+    title: Optional[str],
+    output_path: Optional[str],
+    figsize: Tuple[int, int],
+    show_legend: bool,
+    show_axes: bool,
+    dpi: int,
+    palette: Optional[Dict[str, str]],
+) -> plt.Figure:
+    if metrics is None:
+        metrics = sorted(q_df["metric"].unique().tolist())
+    if not metrics:
+        raise ValueError("No metrics to plot.")
+    if q_df.empty:
+        raise ValueError("The provided DataFrame is empty.")
+
+    variants = sorted(q_df["variant"].unique().tolist())
+    if palette is None:
+        colours: Dict[str, str] = dict(zip(variants, sns.color_palette("colorblind", len(variants))))
+    else:
+        colours = {v: palette.get(v, "#888888") for v in variants}
+
+    n = len(metrics)
+    fig, axes = plt.subplots(1, n, figsize=(figsize[0] * n, figsize[1]))
+    if n == 1:
+        axes = [axes]
+
+    for k, metric in enumerate(metrics):
+        ax = axes[k]
+        sub = q_df[q_df["metric"] == metric]
+        if sub.empty:
+            ax.set_title(metric)
+            continue
+        for variant in variants:
+            vsub = sub[sub["variant"] == variant].sort_values("step")
+            if vsub.empty:
+                continue
+            colour = colours.get(variant)
+            steps = vsub["step"].values
+            q50 = smooth(vsub["q50"].values, smooth_weight)
+            q25 = smooth(vsub["q25"].values, smooth_weight)
+            q75 = smooth(vsub["q75"].values, smooth_weight)
+            ax.plot(steps, q50, color=colour, linewidth=2.0, label=variant)
+            ax.fill_between(steps, q25, q75, color=colour, alpha=0.25, linewidth=0)
+
+        ax.set_title(_pretty_name(metric) if title is None else title)
+        if show_axes:
+            ax.set_ylabel(_pretty_name(metric))
+            ax.set_xlabel("Training steps (in thousands)")
+            ax.xaxis.set_major_formatter(_thousands_formatter())
+        else:
+            ax.set_ylabel("")
+            ax.set_xlabel("")
+        ax.tick_params(axis="both", which="both", top=True, right=True)
+
+    if show_legend:
+        handles, labels = axes[0].get_legend_handles_labels()
+        if handles:
+            fig.legend(
+                handles, labels, loc="upper center",
+                bbox_to_anchor=(0.5, 0.02), ncol=min(len(handles), 6),
+                frameon=True, edgecolor="white", fancybox=False,
+                columnspacing=1.5, handlelength=2.0,
+            )
+
+    fig.tight_layout(rect=[0, 0.08, 1, 1] if show_legend else [0, 0, 1, 1])
+
+    if output_path is not None:
+        os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True)
+        fig.savefig(output_path, bbox_inches="tight", pad_inches=0.15, dpi=dpi)
+
+    return fig
+
+
 _PRETTY: Dict[str, str] = {
     "ep_reward": "Reward",
     "ep_length": "Episode Length",