[Test] Add off-policy data-parallel learning suite

sota-implementations/offpolicy-dp/README.md

@@ -0,0 +1,144 @@
+# Off-policy data-parallel validation
+
+This suite validates TorchRL's Ray-owned replay, asynchronous collection,
+multi-rank learner execution, and direct learner-to-collector weight publication
+for DQN, SAC, DDPG, and TD3.
+
+All four algorithms use transition replay. Each inner Ray collector installs a
+postprocessor through `collector_kwargs` that materializes and flattens `[B, T]`
+rollouts to `[B * T]` before direct insertion. The suite asserts that replay
+`write_count` equals the collector's transition count and that sampled batches
+have shape `[N]`. Grouped sequence replay is not inferred from rollout shape;
+sequence-aware workloads must configure and validate it explicitly.
+
+Two asynchronous evaluators periodically load an exact learner-weight snapshot.
+The metric evaluator runs termination-aware deterministic episodes without
+rendering. A separate single-environment diagnostic rollout records
+`evaluation/video` to W&B; for locomotion tasks it continues until the configured
+video horizon after the agent becomes unhealthy, rather than producing a
+one-frame video when an early policy falls immediately. Rendering is kept out
+of the 100,000 training environments. Both evaluators default to
+`mujoco-torch` on `cuda:7`, avoiding the CPU physics bottleneck and unavailable
+headless EGL/OpenGL support.
+
+The continuous environment is selected by `algorithm.environment_name`. The
+`humanoid_smoke` and `humanoid_scale` profiles run the substantially more
+complex Humanoid task while retaining flat transition replay and the same
+Ray/Gloo/NCCL topology.
+
+Collection remains stochastic after replay prefill. SAC samples its policy,
+TD3 and DDPG apply persistent Gaussian action noise, and DQN retains nonzero
+epsilon-greedy exploration. Evaluation always uses the policy without the DDPG
+or TD3 exploration module. Replay samples log action mean, standard deviation,
+absolute maximum, and saturation fraction so a deterministic or saturated
+collector policy is visible during a run.
+
+The scale profile keeps Ray as the replay owner and uses the stack's distributed
+Gloo tensor transport for collector inserts and learner samples. This avoids
+serializing 25,000-transition GPU-collector batches through Ray pickle while
+retaining the same replay service and ownership topology.
+
+The continuous scale profile uses four GPU collector actors. Each actor owns a
+compiled `mujoco-torch` Hopper environment with 25,000 parallel environments.
+Four additional GPU actors form the NCCL learner group, giving 100,000 total
+environments and four learner ranks on an eight-GPU node.
+
+The learning-scale profiles collect 200 million transitions, or 2,000 policy
+decisions per environment. They do not use a separate random-action phase.
+Instead, the stochastic collection policy fills replay with ten million
+transitions (100 decisions per environment) before optimization starts. This
+separates learner startup from action selection and avoids distributed random
+warmup accounting against a shared replay write counter. Requiring a full
+1,000-step trajectory from every environment would delay learning until 100
+million transitions, so the prefill deliberately targets early terminations
+while persistent exploration continues throughout training.
+
+Scale replay retains the most recent 20 million transitions. A global batch of
+16,384 gives each learner rank 4,096 samples, and 25 optimizer steps per roughly
+100,000 collected frames preserve a sample update-to-data ratio of 4.096 while
+reducing small-batch synchronization overhead. Metric evaluation uses 64
+episodes of up to 1,000 steps every ten million frames; video is recorded every
+20 million frames.
+
+## Installation
+
+Install the experiment-only dependencies into the active environment:
+
+```bash
+uv pip install --python /root/venv/bin/python \
+  -r sota-implementations/offpolicy-dp/requirements.txt
+```
+
+Configure W&B before launching. Credentials are intentionally not read by the
+scripts and must never be stored in this directory.
+
+## Individual runs
+
+Run the planned 100,000-frame DQN validation:
+
+```bash
+python sota-implementations/offpolicy-dp/train.py \
+  algorithm=dqn profile=dqn
+```
+
+Run a reduced SAC smoke test:
+
+```bash
+python sota-implementations/offpolicy-dp/train.py \
+  algorithm=sac profile=smoke_continuous
+```
+
+Run a 100,000-environment, 200-million-transition continuous experiment by
+choosing `sac`, `ddpg`, or `td3`:
+
+```bash
+python sota-implementations/offpolicy-dp/train.py \
+  algorithm=sac profile=scale
+```
+
+Run rendered Humanoid DDPG validation:
+
+```bash
+python sota-implementations/offpolicy-dp/train.py \
+  algorithm=ddpg profile=humanoid_smoke
+python sota-implementations/offpolicy-dp/train.py \
+  algorithm=ddpg profile=humanoid_scale
+```
+
+The Humanoid profiles use `frame_skip=1` so health termination is checked after
+every physics step. This prevents an unstable state in a 25,000-environment
+collector batch from advancing through a five-substep action before reset. The
+scale run is lengthened to 2,000 decisions per environment to compensate for
+the finer control interval; it is not directly comparable to the standard
+five-substep Humanoid benchmark.
+
+Hydra overrides can reduce or expand any resource or training setting. Runtime
+summaries and Hydra output are written under `/root/artifacts/offpolicy-dp`.
+
+## Suite launcher
+
+The launcher uses a fresh Python process for every algorithm so Ray resources
+are released between runs. It records all failures instead of stopping at the
+first one.
+
+```bash
+python sota-implementations/offpolicy-dp/run_suite.py --mode smoke
+python sota-implementations/offpolicy-dp/run_suite.py --mode full
+```
+
+The `all` mode runs reduced smokes first and then the planned full runs.
+
+## Success checks
+
+Every run fails if optimization never starts, replay writes stop early, a
+metric becomes non-finite, the learner's published model version differs from
+its optimization count, or sampled replay data never observes an updated
+policy version. W&B additionally records collection and optimization
+throughput, replay activity, policy-version statistics, and terminal returns
+sampled from replay. Continuous runs also record replay action dispersion and
+saturation to verify that the collection policy remains stochastic.
+
+Learner metrics are logged with slash-delimited namespaces. W&B therefore
+places losses under `loss`, predicted and target values under `value`, gradient
+and step metrics under `optimization`, and entropy/temperature under `policy`,
+rather than leaving these histories in the generic Charts section.

sota-implementations/offpolicy-dp/config/algorithm/ddpg.yaml

@@ -0,0 +1,10 @@
+name: ddpg
+environment_kind: mujoco
+environment_name: hopper
+learning_rate: 0.0003
+gamma: 0.99
+hidden_sizes: [256, 256]
+target_update_polyak: 0.995
+# Keep collection stochastic throughout training. Evaluation uses the actor
+# without this module and remains deterministic.
+exploration_std: 0.2

sota-implementations/offpolicy-dp/config/algorithm/dqn.yaml

@@ -0,0 +1,10 @@
+name: dqn
+environment_kind: gym
+environment_name: CartPole-v1
+learning_rate: 0.0003
+gamma: 0.99
+hidden_sizes: [256, 256]
+target_update_interval: 100
+eps_init: 1.0
+eps_end: 0.05
+annealing_num_steps: 100_000

sota-implementations/offpolicy-dp/config/algorithm/sac.yaml

@@ -0,0 +1,8 @@
+name: sac
+environment_kind: mujoco
+environment_name: hopper
+learning_rate: 0.0003
+gamma: 0.99
+hidden_sizes: [256, 256]
+alpha_init: 1.0
+target_update_polyak: 0.995

sota-implementations/offpolicy-dp/config/algorithm/td3.yaml

@@ -0,0 +1,11 @@
+name: td3
+environment_kind: mujoco
+environment_name: hopper
+learning_rate: 0.0003
+gamma: 0.99
+hidden_sizes: [256, 256]
+target_update_polyak: 0.995
+policy_noise: 0.2
+noise_clip: 0.5
+exploration_std: 0.1
+policy_update_delay: 2

sota-implementations/offpolicy-dp/config/config.yaml

@@ -0,0 +1,64 @@
+defaults:
+  - _self_
+  - algorithm: sac
+  - profile: scale
+
+seed: 42
+
+environment:
+  # Use the task default unless a stability profile requests finer-grained
+  # termination checks between physics substeps.
+  frame_skip: null
+
+logging:
+  project: torchrl-dp-offpolicy-validation
+  group: dp-stack-main
+  mode: online
+  interval_frames: 100_000
+  sample_size: 8192
+  log_env_packages: true
+
+evaluation:
+  enabled: true
+  # Rendering stays in a single environment on one already-allocated device.
+  backend: mujoco-torch
+  device: cuda:7
+  compile_step: false
+  frame_skip: null
+  interval_frames: 2_000_000
+  num_envs: 16
+  num_trajectories: 16
+  max_steps: 100
+  shutdown_timeout: 600.0
+  video: true
+  # Keep scalar evaluation benchmark-correct. Video uses a single diagnostic
+  # rollout that continues after Hopper becomes unhealthy so it shows motion.
+  video_interval_frames: 2_000_000
+  video_num_envs: 1
+  video_num_trajectories: 1
+  video_max_steps: 100
+  render_width: 128
+  render_height: 128
+  video_skip: 2
+  video_fps: 15
+  video_max_frames: 50
+
+artifacts_dir: /root/artifacts/offpolicy-dp
+
+replay:
+  # These feed-forward off-policy algorithms store individual transitions.
+  # A sequence-aware workload must opt into and validate grouped storage itself.
+  storage_mode: transitions
+  transport_backend: gloo
+  transport_timeout: 1200.0
+
+ray:
+  include_dashboard: false
+  log_to_driver: true
+  ignore_reinit_error: true
+
+hydra:
+  job:
+    chdir: false
+  run:
+    dir: /root/artifacts/offpolicy-dp/hydra/${algorithm.name}-${now:%Y%m%d-%H%M%S}

sota-implementations/offpolicy-dp/config/profile/dqn.yaml

@@ -0,0 +1,44 @@
+# @package _global_
+
+environment:
+  backend: gym
+  total_num_envs: 2
+  num_collectors: 2
+  compile_step: false
+  compile_mode: null
+  max_episode_steps: 500
+
+collection:
+  total_frames: 100_000
+  # Epsilon starts at one, so replay prefill is already random without a
+  # separate collector-level random-action phase.
+  init_random_frames: 0
+  learning_starts: 10_000
+  frames_per_env: 1_024
+
+replay:
+  capacity: 100_000
+  batch_size: 512
+  num_cpus: 1
+  transport: auto
+
+learner:
+  world_size: 2
+  num_cpus_per_rank: 1
+  num_gpus_per_rank: 1
+  backend: nccl
+  optim_steps_per_batch: 1
+  poll_interval: 0.05
+  setup_timeout: 300.0
+  command_timeout: 600.0
+
+logging:
+  interval_frames: 10_000
+
+evaluation:
+  device: cpu
+  num_envs: 1
+  num_trajectories: 1
+  interval_frames: 25_000
+  max_steps: 500
+  video_max_frames: 250

sota-implementations/offpolicy-dp/config/profile/humanoid_scale.yaml

@@ -0,0 +1,58 @@
+# @package _global_
+
+algorithm:
+  environment_name: humanoid
+
+environment:
+  backend: mujoco-torch
+  total_num_envs: 100_000
+  num_collectors: 4
+  compile_step: true
+  compile_mode: default
+  # Check health and reset after every physics step. With 25K environments per
+  # collector, a five-substep action eventually exposes unstable solver state.
+  frame_skip: 1
+  max_episode_steps: 1_000
+
+collection:
+  # At 100K environments this provides 2K decisions per environment instead
+  # of the 200-decision transport validation run.
+  total_frames: 200_000_000
+  init_random_frames: 0
+  # Collect 100 noisy-policy decisions per environment before learning. A full
+  # 1K-step prefill would cost 100M transitions before the first update.
+  learning_starts: 10_000_000
+  frames_per_env: 1
+
+replay:
+  capacity: 20_000_000
+  batch_size: 16_384
+  num_cpus: 2
+  transport: distributed
+
+learner:
+  world_size: 4
+  num_cpus_per_rank: 1
+  num_gpus_per_rank: 1
+  backend: nccl
+  # 25 * 16384 / 100000 = 4.096 sampled transitions per collected transition.
+  optim_steps_per_batch: 25
+  poll_interval: 0.05
+  setup_timeout: 600.0
+  command_timeout: 1_200.0
+
+logging:
+  interval_frames: 1_000_000
+  sample_size: 4_096
+
+evaluation:
+  frame_skip: 1
+  interval_frames: 10_000_000
+  num_envs: 64
+  num_trajectories: 64
+  max_steps: 1_000
+  video_interval_frames: 20_000_000
+  video_num_envs: 1
+  video_num_trajectories: 1
+  video_max_steps: 1_000
+  video_max_frames: 500