One shot the repo

EXTENSIONS.md

@@ -0,0 +1,112 @@
+# Project extensions (16-831)
+
+Additions to FlashSAC for the course project. With the default config the training behaviour is unchanged.
+
+## Config keys
+
+All keys live in `configs/agent/flashSAC.yaml` and are set with `--overrides agent.<key>=<value>`.
+
+| Key | Values (default first) | Meaning |
+|---|---|---|
+| `sampler` | `uniform`, `swd`, `laber` | Replay sampling method |
+| `swd_decay_steps` | `2000` | SWD decay horizon, in interaction steps |
+| `swd_min_weight` | `0.1` | SWD weight floor |
+| `laber_multiplier` | `4` | LaBER candidate batch = multiplier × `sample_batch_size` |
+| `actor_delay` | `fixed`, `ema` | Actor delay rule |
+| `actor_update_period` | `2` | Delay `d` for `actor_delay=fixed` |
+| `actor_delay_ema_interval` | `100` | `k`: critic updates between delay recomputations |
+| `actor_delay_ema_beta` | `0.9` | EMA decay, applied once per interval |
+| `actor_delay_ema_scale` | `2.0` | Delay when the TD error equals its EMA |
+| `actor_delay_ema_min` / `_max` | `1.0` / `8.0` | Clip range of the delay |
+| `actor_lr_schedule` | `actor_step`, `global_step` | What drives the actor/temperature learning-rate schedule |
+
+`log_csv` (default `true`, in `configs/flashSAC_base.yaml`) turns the CSV log on or off.
+
+## Methods
+
+**SWD** (`flash_rl/buffers/torch_buffer.py`). WarpSAC's Sample Weight Decay: a transition added `age` interaction
+steps ago is sampled with probability ∝ `max(swd_min_weight, 1 - age / swd_decay_steps)`. No importance correction,
+as in WarpSAC: both the loss and the BatchNorm statistics follow the recency-biased batch.
+This is FlashSAC + SWD, not full WarpSAC (which also changes the critic).
+
+**LaBER** (`FlashSACAgent._select_laber_batch`). Draw `laber_multiplier × B` transitions uniformly, score each with
+
+    surprise_i = mean over the 2 critics of ‖softmax(logits(s_i, a_i)) − target_probs_i‖₂
+
+and keep `B` of them with probability ∝ surprise (with replacement). The importance weights
+`w_i = mean(surprise) / surprise_i` correct the critic update for this selection in two places: the loss terms, and the
+BatchNorm batch statistics of the critic and target critic (weighted mean and variance). Both corrections are
+self-normalised estimates of their uniform-sampling counterparts, so the update is approximately, not exactly,
+unbiased. Scoring runs with BatchNorm in eval mode. The actor is updated on a uniform slice of the candidates, not on
+the high-surprise selection.
+
+**EMA actor delay** (`flash_rl/agents/flashSAC/actor_delay.py`). Every `interval` critic updates, with `e` the mean
+absolute TD error over those updates:
+
+    ema   = beta * ema + (1 - beta) * e
+    delay = clip(scale * e / ema, min, max)
+
+The delay is real-valued: each critic update earns `1/delay` credit and an actor update costs 1.
+
+**Actor LR schedule.** In FlashSAC the actor's cosine decay advances once per actor update, so its learning-rate
+schedule depends on the delay. `actor_lr_schedule=global_step` ties it to the number of critic updates instead. Use it
+for all runs of a delay comparison (it also changes the `d=2` baseline).
+
+## Running
+
+    scripts/run_humanoid.sh <exp_name> [agent overrides...]
+
+runs the four IsaacLab humanoid tasks for each seed; see the examples at the top of the script.
+
+## Logs
+
+Each run writes `logs/<group_name>/<exp_name>/<env_name>/seed<seed>_<time>/`:
+
+- `config.yaml`: the resolved config of the run.
+- `metrics.csv`: one row per logging interval (1% of training). Metrics are averaged over the interval unless noted.
+  A cell is empty when the metric was not produced in that interval.
+
+| Columns | Meaning |
+|---|---|
+| `env_step`, `wall_time` | Row index: environment steps, and seconds since the logger was created |
+| `avg_return`, `avg_length`, `avg_success_*` | Evaluation results (10 evaluations per run, plus one at step 0) |
+| `eval_env_step`, `eval_train_time`, `eval_num_updates` | Exactly when that evaluation ran. Use these for return curves |
+| `time/train_time` | Seconds of training, excluding evaluation and video recording |
+| `time/collect_time`, `time/update_time`, `time/eval_time` | Cumulative seconds in env stepping, gradient updates, evaluation |
+| `time/env_steps_per_sec`, `time/gpu_max_memory_gb` | Throughput and peak GPU memory |
+| `train/num_updates`, `train/num_actor_updates` | Cumulative critic and actor updates |
+| `critic/td_abs` | Mean abs. TD error between distribution means (importance-weighted under LaBER). Drives the EMA delay |
+| `critic/surprise` | Mean surprise, same weighting |
+| `critic/loss`, `critic/q_mean`, `critic/learning_rate` | Critic loss, mean Q of the batch, learning rate |
+| `actor/loss`, `actor/entropy`, `actor/learning_rate`, `temperature/value` | Actor and temperature state |
+| `replay/age_mean`, `_median`, `_p90`, `_max` | Staleness of the critic's batch, in interaction steps |
+| `replay/batch_td_abs`, `replay/batch_surprise` | Unweighted error of the critic's batch |
+| `replay/age_surprise_corr` | Pearson correlation between age and surprise in the critic's batch |
+| `replay/size` | Transitions in the buffer |
+| `actor_delay/delay`, `td_ratio`, `td_ema` | Current delay, `e / ema`, and the EMA |
+| `actor_delay/actor_updated` | Fraction of critic updates that also updated the actor |
+| `laber/candidate_*` | Surprise, TD error, age and their correlation over the uniform candidates |
+| `laber/effective_sample_frac`, `weight_max`, `unique_frac` | How concentrated the LaBER selection is |
+
+## Videos
+
+Checkpoints are saved at the end of training to `models/<group_name>/<exp_name>/<env_name>/seed<seed>-<time>/step<N>/`.
+To record a trained policy to an mp4 without a display:
+
+    uv run python play_isaaclab.py \
+        --checkpoint_path 'models/humanoid/<exp_name>/Isaac-Velocity-Flat-G1-v0/seed0-<time>/step<N>' \
+        --video --video_length 500 --follow_robot \
+        --overrides env=isaaclab \
+        --overrides env.env_name='Isaac-Velocity-Flat-G1-v0' \
+        --overrides agent.buffer_max_length=1
+
+The video (`rl-video-step-0.mp4`, 50 steps per second) goes to `<checkpoint_path>/videos` unless `--video_dir` is set.
+`--follow_robot` keeps the camera on the robot of environment 0; without it the camera is fixed. Pass any overrides
+that change the network (e.g. `agent.critic_hidden_dim`) exactly as in training. Without `--video` the script opens
+the Isaac Sim viewport as before.
+
+## Tests
+
+    python -m pytest tests/test_extensions.py
+
+CPU-only; covers the SWD sampling distribution, ages, the delay rules, LaBER weights, the LR schedule and the CSV log.

configs/agent/flashSAC.yaml

@@ -12,6 +12,17 @@ buffer_min_length: 10_000
 buffer_device_type: 'cuda'
 sample_batch_size: 2048
 
+# Replay sampling
+#   uniform: FlashSAC default.
+#   swd:     WarpSAC's Sample Weight Decay, p(i) ∝ max(swd_min_weight, 1 - age(i) / swd_decay_steps),
+#            with age in interaction steps.
+#   laber:   draw laber_multiplier * sample_batch_size transitions uniformly, then keep sample_batch_size
+#            of them with probability ∝ TD surprise (importance-weighted critic loss).
+sampler: 'uniform'        # ['uniform', 'swd', 'laber']
+swd_decay_steps: 2000
+swd_min_weight: 0.1
+laber_multiplier: 4
+
 normalize_reward: true
 normalized_G_max: 5.0
 
@@ -32,7 +43,22 @@ actor_hidden_dim: 128
 actor_bc_alpha: 0.0
 actor_noise_zeta_mu: 2.0
 actor_noise_zeta_max: 16
+
+# Actor delay = number of critic updates per actor update.
+#   fixed: delay = actor_update_period (FlashSAC default: 2).
+#   ema:   every `interval` critic updates, with e = mean TD error over those updates,
+#          delay = clip(scale * e / EMA_beta(e), min, max).
+actor_delay: 'fixed'      # ['fixed', 'ema']
 actor_update_period: 2
+actor_delay_ema_interval: 100
+actor_delay_ema_beta: 0.9
+actor_delay_ema_scale: 2.0
+actor_delay_ema_min: 1.0
+actor_delay_ema_max: 8.0
+# What drives the actor/temperature learning-rate schedule.
+#   actor_step:  number of actor updates (FlashSAC default; the schedule then depends on the actor delay).
+#   global_step: number of critic updates (same schedule for every actor delay).
+actor_lr_schedule: 'actor_step'   # ['actor_step', 'global_step']
 
 critic_num_blocks: 2
 critic_hidden_dim: 256

configs/flashSAC_base.yaml

@@ -9,6 +9,7 @@ exp_name: 'test'
 seed: 0
 
 logger_type: 'tensorboard'  # ['wandb','tensorboard']
+log_csv: true               # also write logs/<group>/<exp>/<env>/seed<seed>_<time>/{metrics.csv,config.yaml}
 save_path: 'models/${group_name}/${exp_name}/${env.env_name}/seed${seed}-TIMESTAMP'
 agent_load_path: null
 buffer_load_path: null

flash_rl/agents/flashSAC/actor_delay.py

@@ -0,0 +1,76 @@
+from typing import Any, Optional
+
+
+class ActorDelay:
+    """
+    Decides, at every critic update, whether the actor is updated too.
+    `delay` is the number of critic updates per actor update.
+
+    fixed: delay = period (FlashSAC uses 2).
+    ema:   every `interval` critic updates, with e = mean TD error over those updates,
+               ema   = beta * ema + (1 - beta) * e
+               delay = clip(scale * e / ema, min_delay, max_delay)
+           The actor is updated less often while the critic error is above its recent average.
+    """
+
+    def __init__(
+        self,
+        mode: str,
+        period: int,
+        interval: int,
+        beta: float,
+        scale: float,
+        min_delay: float,
+        max_delay: float,
+    ):
+        assert mode in ("fixed", "ema"), f"Unknown actor delay: {mode}"
+        assert interval >= 1 and 0.0 <= beta < 1.0 and 0.0 < min_delay <= max_delay
+        self._mode = mode
+        self._interval = interval
+        self._beta = beta
+        self._scale = scale
+        self._min_delay = min_delay
+        self._max_delay = max_delay
+
+        # until the first interval ends, the error ratio is taken to be 1
+        self.delay = float(period) if mode == "fixed" else self._clip(scale)
+        self.ratio = 1.0
+        self.ema: Optional[float] = None
+        self._window_sum = 0.0
+        self._window_count = 0
+        # `delay` is real-valued: each critic update earns 1/delay credit, an actor update costs 1
+        self._credit = 1.0
+
+    def _clip(self, delay: float) -> float:
+        return min(max(delay, self._min_delay), self._max_delay)
+
+    def should_update_actor(self, update_step: int) -> bool:
+        if self._mode == "fixed":
+            return update_step % int(self.delay) == 0
+        do_update = self._credit >= 1.0 - 1e-6
+        if do_update:
+            self._credit -= 1.0
+        self._credit += 1.0 / self.delay
+        return do_update
+
+    def observe(self, td_error: float) -> None:
+        """Record the TD error of one critic update."""
+        if self._mode == "fixed":
+            return
+        self._window_sum += td_error
+        self._window_count += 1
+        if self._window_count < self._interval:
+            return
+        error = self._window_sum / self._window_count
+        self._window_sum, self._window_count = 0.0, 0
+        self.ema = error if self.ema is None else self._beta * self.ema + (1 - self._beta) * error
+        self.ratio = error / (self.ema + 1e-8)
+        self.delay = self._clip(self._scale * self.ratio)
+
+    def state_dict(self) -> dict[str, Any]:
+        keys = ("delay", "ratio", "ema", "_window_sum", "_window_count", "_credit")
+        return {key: getattr(self, key) for key in keys}
+
+    def load_state_dict(self, state: dict[str, Any]) -> None:
+        for key, value in state.items():
+            setattr(self, key, value)