Tabular Safe RL Baselines

This tutorial compares MASA’s tabular safe RL baselines on one constrained tabular environment. The point is to read the reward/safety tradeoff each method encodes, not to rank algorithms from a tiny run.

Runnable notebook: notebooks/tutorials/06_tabular_safe_rl_baselines.ipynb

Shared Setup

Use CPU-first setup before importing MASA/JAX modules:

import os

os.environ.setdefault("JAX_PLATFORMS", "cpu")
os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "2")

All algorithms use the same constrained environment:

from masa.common.utils import make_env
from masa.envs.tabular.colour_grid_world import cost_fn, label_fn

def make_train_env():
    return make_env(
        "colour_grid_world",
        "cmdp",
        40,
        label_fn=label_fn,
        cost_fn=cost_fn,
        budget=0.0,
    )

colour_grid_world is small enough for tutorial execution, has discrete states/actions, and separates sparse reward from safety cost:

• goal gives task reward,

• blue gives safety cost,

• the CMDP budget is 0.0.

Probe the Signals

Before training, run scripted rollouts to inspect the data that algorithms receive:

• unsafe script: seed 1, actions [2, 2, 2, 2], reaches blue,

• goal script: seed 4, actions [2] * 8 + [1] * 8, reaches goal.

The unsafe script shows cost=1.0, violation=1.0, cum_cost > 0, and satisfied=0.0. The goal script shows reward without violating the CMDP budget.

Algorithms Compared

Algorithm	Class	Safety mechanism
q_learning	QL	Task Q table only; costs are logged but not penalized in the update target.
q_learning_lambda	QL_Lambda	Subtracts cost_lambda * cost from the reward target.
lcrl	LCRL	Maps violations to an absorbing-style value based on r_min / (1 - gamma).
sem	SEM	Learns task Q plus auxiliary D and C safety tables that alter action selection.
recreg	RECREG	Learns task Q, backup B, risk S, and can report override_rate when actions are replaced.

The tutorial uses this compact registry:

ALGORITHMS = {
    "q_learning": (QL, {}),
    "q_learning_lambda": (QL_Lambda, {"cost_lambda": 1.0}),
    "lcrl": (LCRL, {"r_min": -1.0}),
    "sem": (SEM, {"r_min": -1.0, "cost_coef": 1.0}),
    "recreg": (
        RECREG,
        {
            "mode": "model_free",
            "model_checking": "none",
            "horizon": 3,
            "safety_prob": 0.2,
            "cost_coef": 2.0,
        },
    ),
}

Tiny Smoke Runs

Each algorithm is trained with the same tiny configuration:

algo.train(
    num_frames=20,
    eval_freq=10,
    log_freq=10,
    num_eval_episodes=1,
    stats_window_size=5,
)

These runs are intentionally too small for performance claims. They are diagnostics that prove:

• all five algorithms can train on the same MASA wrapper stack,

• train/eval logging includes reward and constraint metrics,

• the learned objects expose the expected state tables.

Reading the Learned State

After training, inspect the object shapes:

• QL, QL_Lambda, and LCRL expose Q,

• SEM exposes Q, D, and C,

• RECREG exposes Q, backup table B, and risk table S.

Use the printed train/rollout, train/stats, and eval/rollout blocks to interpret reward and safety together. For example, RECREG may add override_rate to episode metrics when it replaces a risky task action with a backup action.

What to Take Away

The algorithms differ less in the environment interaction loop than in how they turn cost and violation signals into learning or action-selection pressure:

• QL is the task-only baseline,

• QL_Lambda is a soft penalty baseline,

• LCRL makes violations strongly unattractive,

• SEM separates task and safety estimates,

• RECREG introduces backup-policy intervention.

For real comparisons, increase training frames, run multiple seeds, and summarize confidence intervals. Keep this tutorial run as a quick wiring and interpretation check.