diff --git a/.github/workflows/unit-tests.yml b/.github/workflows/unit-tests.yml new file mode 100644 index 0000000..9a0fc52 --- /dev/null +++ b/.github/workflows/unit-tests.yml @@ -0,0 +1,32 @@ +name: Unit Tests + +on: + push: + branches: + - main + - master + pull_request: + +jobs: + unit-tests: + runs-on: ubuntu-latest + + steps: + - name: Checkout + uses: actions/checkout@v4 + + - name: Setup Python + uses: actions/setup-python@v5 + with: + python-version: "3.11" + + - name: Setup uv + uses: astral-sh/setup-uv@v5 + + - name: Install project and test dependencies + run: | + uv venv .venv + uv pip install --python .venv/bin/python --torch-backend cpu -e '.[dev]' + + - name: Run unit tests + run: .venv/bin/python -m pytest tests/unit diff --git a/.gitignore b/.gitignore index 8476755..1dae27a 100644 --- a/.gitignore +++ b/.gitignore @@ -40,6 +40,7 @@ pip-delete-this-directory.txt htmlcov/ .tox/ .nox/ +.pytest_artifacts/ .coverage .coverage.* .cache @@ -140,6 +141,9 @@ logs/ runs/ outputs/ output/ +.o/ +.rt/ +.w/ # runs resource_pool_auto.yaml diff --git a/examples/wbc_tracking/pyproject.toml b/examples/wbc_tracking/pyproject.toml index b977df3..c617646 100644 --- a/examples/wbc_tracking/pyproject.toml +++ b/examples/wbc_tracking/pyproject.toml @@ -15,9 +15,18 @@ override-dependencies = [ "sympy==1.13.1" ] +[[tool.uv.index]] +name = "pypi" +url = "https://pypi.org/simple" + +[[tool.uv.index]] +name = "nvidia" +url = "https://pypi.nvidia.com" + [tool.uv.sources] rlightning = {path = "../../", editable = true} rsl-rl = {path = "../../third_party/rsl_rl", editable = true} +isaaclab = { index = "nvidia" } [tool.uv.extra-build-dependencies] flatdict = ["setuptools<81"] diff --git a/pyproject.toml b/pyproject.toml index bb55592..a11634e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,3 +1,7 @@ +[build-system] +requires = ["setuptools>=80.8.0,<81", "wheel"] +build-backend = "setuptools.build_meta" + [project] name = "rlightning" version = "0.1.0" @@ -26,7 +30,7 @@ dependencies = [ "transformers", "uvloop", "wandb", - "zmq", + "pyzmq", "zstandard>=0.25.0", ] @@ -34,6 +38,8 @@ dependencies = [ dev = [ "debugpy>=1.8.14", "ipython", + "pytest>=8.3.5", + "pytest-cov>=6.1.1", ] isaaclab = [ "isaaclab[isaacsim]==2.2.0", @@ -54,13 +60,17 @@ humanoid = [ "scipy", "easydict", "joblib", - "smplx @ git+https://github.com/ZhengyiLuo/smplx.git@master", - "smpl-sim @ git+https://github.com/ZhengyiLuo/SMPLSim.git@master", "open3d==0.19.0", "natsort==8.4.0", "mink==0.0.13", ] +[dependency-groups] +humanoid-dev = [ + "smplx @ git+https://github.com/ZhengyiLuo/smplx.git@master", + "smpl-sim @ git+https://github.com/ZhengyiLuo/SMPLSim.git@master", +] + [tool.isort] profile = "black" @@ -73,6 +83,19 @@ typeCheckingMode = "off" reportMissingImports = false reportMissingModuleSource = false +[tool.pytest.ini_options] +minversion = "8.0" +testpaths = ["tests"] +addopts = "-ra -m 'not integration and not slow and not gpu and not maniskill and not isaaclab and not e2e'" +markers = [ + "slow: marks tests that are slower or intended for scheduled regression runs", + "gpu: marks tests that require a GPU runtime", + "integration: marks tests that cover multiple components working together", + "e2e: marks end-to-end CLI or workflow tests", + "maniskill: marks tests that require the ManiSkill stack", + "isaaclab: marks tests that require the IsaacLab stack", +] + [tool.setuptools.packages.find] include = ["rlightning*"] diff --git a/rlightning/engine/async_rl_engine.py b/rlightning/engine/async_rl_engine.py index b967b9f..92b6e8a 100644 --- a/rlightning/engine/async_rl_engine.py +++ b/rlightning/engine/async_rl_engine.py @@ -249,7 +249,7 @@ def _train_loop(self) -> None: self.coordinator.wait_for_dataset_ready() self.coordinator.wait_for_weights_updated() self._train() - if self.config.train.save_interval > 0 and (self.epoch + 1) % self.config.train.save_interval == 0: + if self.config.train.save_interval > 0 and self.epoch % self.config.train.save_interval == 0: ckpt_path = f"{self.config.train.save_dir}/epoch_{self.epoch}.pt" self.policy_group.save_checkpoint(path=ckpt_path) self.coordinator.notify_train_step_done() diff --git a/rlightning/engine/sync_rl_engine.py b/rlightning/engine/sync_rl_engine.py index 0dac126..67eb9be 100644 --- a/rlightning/engine/sync_rl_engine.py +++ b/rlightning/engine/sync_rl_engine.py @@ -215,21 +215,19 @@ def run(self) -> None: Executes the training loop for the configured number of epochs, performing rollout, training, and periodic evaluation. """ - logger.info("Evaluating before training...") - self._evaluate(obj_set="train", prefix="eval") - self._evaluate(obj_set="test", prefix="eval_ood") - for self.epoch in self.iter_epochs(num_epochs=self.config.train.max_epochs): self._rollout(obj_set="train", prefix="rollout") self._update_dataset() self._train() - self._sync_weights() - if self.config.train.eval_interval > 0 and (self.epoch + 1) % self.config.train.save_interval == 0: + if self.config.train.eval_interval > 0 and self.epoch % self.config.train.save_interval == 0: ckpt_path = f"{self.config.train.save_dir}/epoch_{self.epoch}.pt" self.policy_group.save_checkpoint(path=ckpt_path) + + # sync weights after training and save checkpoint + self._sync_weights() - if self.config.train.eval_interval > 0 and (self.epoch + 1) % self.config.train.eval_interval == 0: + if self.config.train.eval_interval > 0 and self.epoch % self.config.train.eval_interval == 0: logger.info(f"Evaluating at epoch {self.epoch}") self._evaluate(obj_set="train", prefix="eval") self._evaluate(obj_set="test", prefix="eval_ood") diff --git a/rlightning/policy/base_policy.py b/rlightning/policy/base_policy.py index 09b6b66..7b7f8dd 100644 --- a/rlightning/policy/base_policy.py +++ b/rlightning/policy/base_policy.py @@ -60,6 +60,23 @@ def infer_train_dp_world_size() -> int: return 1 +def clone_checkpoint_value(value: Any) -> Any: + """Recursively clone checkpoint payloads into CPU-owned tensors.""" + if torch.is_tensor(value): + return value.detach().cpu().clone() + + if isinstance(value, dict): + return value.__class__((key, clone_checkpoint_value(child)) for key, child in value.items()) + + if isinstance(value, list): + return [clone_checkpoint_value(child) for child in value] + + if isinstance(value, tuple): + return tuple(clone_checkpoint_value(child) for child in value) + + return value + + class PolicyRole(StrEnum): """Policy role enumeration.""" @@ -548,15 +565,20 @@ def save_checkpoint(self, path: str) -> None: ckpt_folder = Path(path).parent os.makedirs(ckpt_folder, exist_ok=True) - state: Dict[str, Dict] = {} - for name, model in self.model_list: - if isinstance(model, DDP): - module = model.module - else: - module = model - state[name] = module.state_dict() + model_was_offloaded = getattr(self, "_model_params_offloaded", False) + if model_was_offloaded: + self.reload_model_param_and_grad(load_grad=False) - torch.save(state, path) + try: + state: Dict[str, Dict] = {} + for name, model in self.model_list: + module = model.module if isinstance(model, DDP) else model + state[name] = clone_checkpoint_value(module.state_dict()) + + torch.save(state, path) + finally: + if model_was_offloaded: + self.offload_model_param_and_grad(offload_grad=False) def reset_training_state( self, train_config: TrainConfig, env_meta: Optional[Any] = None, seed: Optional[int] = None diff --git a/rlightning/weights/weight_buffer_mixin.py b/rlightning/weights/weight_buffer_mixin.py index f0ebc08..5753e55 100644 --- a/rlightning/weights/weight_buffer_mixin.py +++ b/rlightning/weights/weight_buffer_mixin.py @@ -43,6 +43,7 @@ def __init_weight_buffer_mixin__(self, buffer_strategy: str): # for offload model param and grad self.cpu_param_backup = {} + self._model_params_offloaded = False def init_weight_buffer(self, shared_weight_buffer=None): """Initialize the weight buffer.""" @@ -177,10 +178,11 @@ def offload_model_param_and_grad(self, offload_grad=False): actual_model = self.model.module if isinstance(self.model, DDP) else self.model for name, param in actual_model.named_parameters(): if param.data.storage().size() > 0: - self.cpu_param_backup[name] = (param.data.cpu(), param.data.size()) + self.cpu_param_backup[name] = (param.data.detach().cpu().clone(), param.data.size()) _free_storage(param.data) if offload_grad and param.grad is not None: param.grad = param.grad.to("cpu", non_blocking=True) + self._model_params_offloaded = True self.clear_memory(sync=True) profiler.log_gpu_memory_usage("offload_model_param_and_grad") @@ -198,6 +200,7 @@ def reload_model_param_and_grad(self, load_grad=False): if load_grad and param.grad is not None: param.grad = param.grad.to(self.device, non_blocking=True) + self._model_params_offloaded = False self.clear_memory(sync=True) profiler.log_gpu_memory_usage("reload_model_param_and_grad") diff --git a/scripts/testing/run_all_tests.sh b/scripts/testing/run_all_tests.sh new file mode 100755 index 0000000..efb0dd9 --- /dev/null +++ b/scripts/testing/run_all_tests.sh @@ -0,0 +1,227 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd) + +usage() { + cat <<'EOF' +Usage: bash scripts/testing/run_all_tests.sh [options] [targets...] + +Runs the local test workflow. With no targets, this runs: + 1. unit + 2. integration core + 3. integration isaaclab + 4. e2e + +Targets: + unit Run tests/unit in the main .venv. + integration Run both integration profiles: core + isaaclab. + integration-core Run only the core integration profile. + integration-isaaclab Run only the isaaclab integration profile. + e2e Run tests/e2e in the main .venv. + all Run the full workflow. This is the default. + +Options: + -q, --quiet Use concise pytest output. This is the default. + -v, --verbose Use verbose pytest output and print live E2E logs. + --no-setup Do not auto-create missing integration environments. + -h, --help Show this help message. + +Examples: + bash scripts/testing/run_all_tests.sh + bash scripts/testing/run_all_tests.sh --verbose + bash scripts/testing/run_all_tests.sh unit + bash scripts/testing/run_all_tests.sh integration + bash scripts/testing/run_all_tests.sh integration-core integration-isaaclab + bash scripts/testing/run_all_tests.sh e2e +EOF +} + +verbosity="quiet" +auto_setup="1" +targets=() +resolved_targets=() + +while [[ $# -gt 0 ]]; do + case "$1" in + -q|--quiet) + verbosity="quiet" + shift + ;; + -v|--verbose) + verbosity="verbose" + shift + ;; + --no-setup) + auto_setup="0" + shift + ;; + -h|--help) + usage + exit 0 + ;; + unit|integration|integration-core|integration-isaaclab|e2e|all) + targets+=("$1") + shift + ;; + *) + echo "Unknown option or target: $1" >&2 + usage + exit 1 + ;; + esac +done + +log_section() { + echo + echo "==> $1" +} + +ensure_main_env() { + if [[ ! -x "$ROOT_DIR/.venv/bin/python" ]]; then + echo "Missing main test environment: $ROOT_DIR/.venv" >&2 + echo "Prepare it first, for example:" >&2 + echo " uv venv .venv" >&2 + echo " uv pip install --python .venv/bin/python --torch-backend cpu -e '.[dev]'" >&2 + exit 1 + fi +} + +ensure_integration_env() { + local profile="$1" + local venv_dir="$ROOT_DIR/.venv-int-$profile" + + if [[ -x "$venv_dir/bin/python" ]]; then + return + fi + + if [[ "$auto_setup" != "1" ]]; then + echo "Missing integration environment: $venv_dir" >&2 + echo "Run: bash scripts/testing/setup_integration_env.sh $profile" >&2 + exit 1 + fi + + log_section "Preparing integration environment: $profile" + bash "$ROOT_DIR/scripts/testing/setup_integration_env.sh" "$profile" +} + +append_target_once() { + local target="$1" + local existing + for existing in "${resolved_targets[@]}"; do + if [[ "$existing" == "$target" ]]; then + return + fi + done + resolved_targets+=("$target") +} + +resolve_targets() { + local requested_targets=("${targets[@]}") + + if [[ ${#requested_targets[@]} -eq 0 ]]; then + requested_targets=(all) + fi + + local target + for target in "${requested_targets[@]}"; do + case "$target" in + all) + append_target_once unit + append_target_once integration-core + append_target_once integration-isaaclab + append_target_once e2e + ;; + integration) + append_target_once integration-core + append_target_once integration-isaaclab + ;; + unit|integration-core|integration-isaaclab|e2e) + append_target_once "$target" + ;; + *) + echo "Unknown target: $target" >&2 + usage + exit 1 + ;; + esac + done +} + +run_unit_tests() { + log_section "Unit tests" + if [[ "$verbosity" == "verbose" ]]; then + "$ROOT_DIR/.venv/bin/python" -m pytest -vv -s -ra tests/unit + else + "$ROOT_DIR/.venv/bin/python" -m pytest -q tests/unit + fi +} + +run_integration_tests() { + local profile="$1" + log_section "Integration tests ($profile)" + if [[ "$verbosity" == "verbose" ]]; then + bash "$ROOT_DIR/scripts/testing/run_integration_tests.sh" "$profile" -vv -s -ra + else + bash "$ROOT_DIR/scripts/testing/run_integration_tests.sh" "$profile" -q + fi +} + +run_e2e_tests() { + log_section "E2E tests" + if [[ "$verbosity" == "verbose" ]]; then + RLIGHTNING_E2E_LIVE_LOGS=1 RAY_DEDUP_LOGS=0 \ + "$ROOT_DIR/.venv/bin/python" -m pytest -o addopts='' -vv -s -ra tests/e2e + else + "$ROOT_DIR/.venv/bin/python" -m pytest -o addopts='' -q tests/e2e + fi +} + +resolve_targets + +need_main_env="0" +need_core_integration="0" +need_isaaclab_integration="0" + +for target in "${resolved_targets[@]}"; do + case "$target" in + unit|e2e) + need_main_env="1" + ;; + integration-core) + need_core_integration="1" + ;; + integration-isaaclab) + need_isaaclab_integration="1" + ;; + esac +done + +if [[ "$need_main_env" == "1" ]]; then + ensure_main_env +fi +if [[ "$need_core_integration" == "1" ]]; then + ensure_integration_env core +fi +if [[ "$need_isaaclab_integration" == "1" ]]; then + ensure_integration_env isaaclab +fi + +for target in "${resolved_targets[@]}"; do + case "$target" in + unit) + run_unit_tests + ;; + integration-core) + run_integration_tests core + ;; + integration-isaaclab) + run_integration_tests isaaclab + ;; + e2e) + run_e2e_tests + ;; + esac +done + +log_section "All tests passed" diff --git a/scripts/testing/run_integration_tests.sh b/scripts/testing/run_integration_tests.sh new file mode 100755 index 0000000..da44af1 --- /dev/null +++ b/scripts/testing/run_integration_tests.sh @@ -0,0 +1,57 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd) + +usage() { + cat <<'EOF' +Usage: bash scripts/testing/run_integration_tests.sh [pytest args...] + +Profiles: + core Run non-IsaacLab integration tests in .venv-int-core. + isaaclab Run IsaacLab integration tests in .venv-int-isaaclab. + +Examples: + bash scripts/testing/run_integration_tests.sh core -q + bash scripts/testing/run_integration_tests.sh core -vv -s -ra + bash scripts/testing/run_integration_tests.sh isaaclab -q + bash scripts/testing/run_integration_tests.sh isaaclab -vv -s -ra +EOF +} + +if [[ "${1:-}" == "-h" || "${1:-}" == "--help" ]]; then + usage + exit 0 +fi + +if [[ $# -lt 1 ]]; then + usage + exit 1 +fi + +profile="$1" +shift + +case "$profile" in + core) + venv_dir="$ROOT_DIR/.venv-int-core" + default_args=(-o addopts='' -m 'not isaaclab' tests/integration) + ;; + isaaclab) + venv_dir="$ROOT_DIR/.venv-int-isaaclab" + default_args=(-o addopts='' -m isaaclab tests/integration/env/test_isaac_manager_based.py) + ;; + *) + echo "Unknown profile: $profile" >&2 + usage + exit 1 + ;; +esac + +if [[ ! -x "$venv_dir/bin/python" ]]; then + echo "Missing virtual environment: $venv_dir" >&2 + echo "Run: bash scripts/testing/setup_integration_env.sh $profile" >&2 + exit 1 +fi + +exec "$venv_dir/bin/python" -m pytest "${default_args[@]}" "$@" diff --git a/scripts/testing/setup_integration_env.sh b/scripts/testing/setup_integration_env.sh new file mode 100755 index 0000000..8080cc3 --- /dev/null +++ b/scripts/testing/setup_integration_env.sh @@ -0,0 +1,60 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd) + +usage() { + cat <<'EOF' +Usage: bash scripts/testing/setup_integration_env.sh + +Profiles: + core Create .venv-int-core with pytest + mujoco + ale extras. + isaaclab Create .venv-int-isaaclab with pytest + isaaclab + humanoid extras and humanoid-dev group. + +Examples: + bash scripts/testing/setup_integration_env.sh core + bash scripts/testing/setup_integration_env.sh isaaclab +EOF +} + +if [[ "${1:-}" == "-h" || "${1:-}" == "--help" ]]; then + usage + exit 0 +fi + +if [[ $# -ne 1 ]]; then + usage + exit 1 +fi + +profile="$1" +extra_args=() + +case "$profile" in + core) + venv_dir="$ROOT_DIR/.venv-int-core" + project_spec='.[mujoco,ale]' + ;; + isaaclab) + venv_dir="$ROOT_DIR/.venv-int-isaaclab" + project_spec='.[isaaclab,humanoid]' + extra_args=(--group humanoid-dev) + ;; + *) + echo "Unknown profile: $profile" >&2 + usage + exit 1 + ;; +esac + +uv venv --allow-existing "$venv_dir" +uv pip install --python "$venv_dir/bin/python" -e "$project_spec" "pytest>=8.3.5" "${extra_args[@]:-}" + +cat < dict: + for key, value in updates.items(): + if isinstance(value, dict) and isinstance(target.get(key), dict): + _deep_update(target[key], value) + else: + target[key] = value + return target + + +@pytest.fixture(autouse=True) +def reset_global_resource_manager() -> None: + GlobalResourceManager.reset() + yield + GlobalResourceManager.reset() + + +@pytest.fixture(autouse=True) +def set_random_seeds(): + set_deterministic_seeds(42) + yield + np.random.seed() + torch.manual_seed(torch.initial_seed()) + random.seed() + + +@pytest.fixture +def event_loop(): + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + yield loop + loop.close() + + +@pytest.fixture +def performance_timer(): + return PerformanceTimer() + + +@pytest.fixture(scope="session") +def ray_multi_node_cluster(): + if should_skip_ray_test() or should_skip_gpu_test(): + yield + return + + env = os.environ.copy() + subprocess.run(["ray", "stop", "--force"], check=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) + + try: + head_env = env.copy() + head_env["CUDA_VISIBLE_DEVICES"] = "0,1" + subprocess.run( + ["ray", "start", "--head", "--disable-usage-stats", "--num-gpus", "2"], + cwd=str(_ROOT), + env=head_env, + check=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + + worker1_env = env.copy() + worker1_env["CUDA_VISIBLE_DEVICES"] = "2,3,4" + subprocess.run( + ["ray", "start", "--address", "auto", "--num-gpus", "3"], + cwd=str(_ROOT), + env=worker1_env, + check=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + + worker2_env = env.copy() + worker2_env["CUDA_VISIBLE_DEVICES"] = "5,6,7" + subprocess.run( + ["ray", "start", "--address", "auto", "--num-gpus", "3"], + cwd=str(_ROOT), + env=worker2_env, + check=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + + yield + finally: + subprocess.run(["ray", "stop", "--force"], check=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) + + +@pytest.fixture +def make_main_config_dict(): + def factory(**overrides): + config = { + "engine": "syncrl", + "env": [ + { + "name": "test-env", + "backend": "mujoco", + "task": "cartpole", + "num_workers": 2, + "num_envs": 1, + "num_cpus": 2, + "num_gpus": 0.0, + } + ], + "buffer": { + "type": "RolloutBuffer", + "capacity": 32, + "storage": { + "type": "unified", + "device": "cpu", + }, + }, + "policy": { + "type": "SimplePPOPolicy", + "train_num_gpus": 1.0, + "eval_num_gpus": 0.5, + "weight_buffer": { + "type": "WeightBuffer", + "buffer_strategy": "Double", + }, + }, + "train": { + "max_epochs": 1, + "batch_size": 8, + }, + "cluster": { + "train_worker_num": 2, + "eval_worker_num": 3, + "buffer_worker_num": 1, + "remote_train": False, + "remote_eval": False, + "remote_storage": False, + "remote_env": False, + }, + } + return _deep_update(copy.deepcopy(config), overrides) + + return factory diff --git a/tests/e2e/__init__.py b/tests/e2e/__init__.py new file mode 100644 index 0000000..98b50e4 --- /dev/null +++ b/tests/e2e/__init__.py @@ -0,0 +1 @@ +"""End-to-end tests.""" diff --git a/tests/e2e/_cli_smoke_utils.py b/tests/e2e/_cli_smoke_utils.py new file mode 100644 index 0000000..bce3ef9 --- /dev/null +++ b/tests/e2e/_cli_smoke_utils.py @@ -0,0 +1,208 @@ +"""Shared helpers for CLI-driven E2E smoke tests.""" + +from __future__ import annotations + +import os +import subprocess +from pathlib import Path + +import pytest +import yaml + + +def build_pythonpath(root: Path) -> str: + parts = [ + str(root), + str(root / "examples"), + str(root / "third_party"), + str(root / "third_party" / "rw_rl"), + str(root / "third_party" / "rw_rl" / "src"), + ] + existing = os.environ.get("PYTHONPATH", "") + return ":".join(parts + ([existing] if existing else [])) + + +def load_yaml(path: Path) -> dict: + with path.open("r", encoding="utf-8") as handle: + return yaml.safe_load(handle) + + +def require_model_path(policy_yaml: Path) -> None: + cfg = load_yaml(policy_yaml) + model_path = cfg["model_cfg"]["model_path"] + if not Path(model_path).exists(): + pytest.skip(f"Model path not found: {model_path}") + + +def example_python(root: Path, example_name: str) -> Path: + python = root / "examples" / example_name / ".venv" / "bin" / "python" + if not python.exists(): + pytest.skip(f"Example virtualenv not found: {python}") + return python + + +def pick_free_gpu_ids(required_count: int, max_used_mb: int = 1024) -> list[str]: + try: + result = subprocess.run( + [ + "nvidia-smi", + "--query-gpu=index,memory.used", + "--format=csv,noheader,nounits", + ], + check=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + except (FileNotFoundError, subprocess.CalledProcessError): + pytest.skip("nvidia-smi is unavailable") + + gpu_rows: list[tuple[int, int]] = [] + for line in result.stdout.splitlines(): + if not line.strip(): + continue + idx_str, used_str = [part.strip() for part in line.split(",", maxsplit=1)] + gpu_rows.append((int(idx_str), int(used_str))) + + free_gpu_ids = [str(idx) for idx, used_mb in sorted(gpu_rows, key=lambda row: row[1]) if used_mb <= max_used_mb] + if len(free_gpu_ids) < required_count: + pytest.skip(f"Need {required_count} mostly idle GPUs (<= {max_used_mb} MiB used), found {free_gpu_ids}") + return free_gpu_ids[:required_count] + + +def resolve_gpu_ids(required_count: int, max_used_mb: int = 1024) -> list[str]: + configured = os.environ.get("CUDA_VISIBLE_DEVICES", "").strip() + if configured: + gpu_ids = [gpu_id.strip() for gpu_id in configured.split(",") if gpu_id.strip()] + if len(gpu_ids) < required_count: + pytest.skip( + f"CUDA_VISIBLE_DEVICES={configured!r} exposes {len(gpu_ids)} GPU(s), " + f"but this test needs {required_count}" + ) + return gpu_ids[:required_count] + return pick_free_gpu_ids(required_count=required_count, max_used_mb=max_used_mb) + + +def short_ray_tmpdir(key: str) -> str: + path = Path("/tmp") / f"r{key}" + path.mkdir(parents=True, exist_ok=True) + return str(path) + + +def run_example( + root: Path, + cmd: list[str], + timeout_s: int, + extra_env: dict[str, str] | None = None, +) -> subprocess.CompletedProcess[str]: + env = os.environ.copy() + env["PYTHONPATH"] = build_pythonpath(root) + env.setdefault("HYDRA_FULL_ERROR", "1") + env.setdefault("PYTHONUNBUFFERED", "1") + env.setdefault("TOKENIZERS_PARALLELISM", "false") + env.setdefault("PYTHONHASHSEED", "0") + if extra_env is not None: + env.update(extra_env) + + live_logs = env.get("RLIGHTNING_E2E_LIVE_LOGS") == "1" + if live_logs: + return subprocess.run( + cmd, + cwd=str(root), + env=env, + check=False, + timeout=timeout_s, + ) + + return subprocess.run( + cmd, + cwd=str(root), + env=env, + check=False, + timeout=timeout_s, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + + +def artifact_mtime(path: Path) -> float | None: + return path.stat().st_mtime if path.exists() else None + + +def assert_updated_artifact(path: Path, previous_mtime: float | None, name: str) -> None: + if not path.exists(): + pytest.fail(f"{name} did not produce expected artifact: {path}") + if previous_mtime is not None and path.stat().st_mtime <= previous_mtime: + pytest.fail(f"{name} did not update expected artifact: {path}") + + +def snapshot_hydra_runs(root: Path) -> set[Path]: + outputs_root = root / "outputs" + if not outputs_root.exists(): + return set() + return {path.resolve() for path in outputs_root.glob("*/*") if path.is_dir()} + + +def find_new_hydra_run_dir(root: Path, before: set[Path], expected_log_name: str) -> Path: + candidates = [] + for path in snapshot_hydra_runs(root) - before: + config_path = path / ".hydra" / "config.yaml" + if not config_path.exists(): + continue + try: + cfg = load_yaml(config_path) + except Exception: + continue + if cfg.get("log", {}).get("name") == expected_log_name: + candidates.append(path) + + if not candidates: + pytest.fail(f"Could not find new Hydra run dir for log.name={expected_log_name}") + return max(candidates, key=lambda path: path.stat().st_mtime) + + +def known_wbc_cleanup_error(output: str) -> bool: + return ( + "ray.exceptions.ActorDiedError" in output + and "env_group.close()" in output + and "IsaacManagerBasedRLEnv" in output + ) + + +def assert_subprocess_success( + result: subprocess.CompletedProcess[str], + name: str, + *, + allow_known_wbc_cleanup_error: bool = False, + expected_paths: list[Path] | None = None, +) -> None: + captured_output = result.stdout is not None or result.stderr is not None + output = (result.stdout or "") + "\n" + (result.stderr or "") if captured_output else "" + if result.returncode == 0: + missing_paths = [path for path in expected_paths or [] if not path.exists()] + if missing_paths: + pytest.fail( + f"{name} exited with code 0 but did not produce expected artifacts: " + + ", ".join(str(path) for path in missing_paths) + ) + if captured_output and "Done." not in output: + pytest.fail(f"{name} finished without the expected completion marker.\n{output[-6000:]}") + return + + if allow_known_wbc_cleanup_error and known_wbc_cleanup_error(output): + missing_paths = [path for path in expected_paths or [] if not path.exists()] + if missing_paths: + pytest.fail( + f"{name} hit the known cleanup error but did not produce expected artifacts: " + + ", ".join(str(path) for path in missing_paths) + ) + return + + out_tail = (result.stdout or "").splitlines()[-80:] + err_tail = (result.stderr or "").splitlines()[-80:] + pytest.fail( + f"{name} failed with return code {result.returncode}.\n" + f"Last stdout lines:\n{chr(10).join(out_tail)}\n\n" + f"Last stderr lines:\n{chr(10).join(err_tail)}" + ) diff --git a/tests/e2e/test_minimal_training_smoke.py b/tests/e2e/test_minimal_training_smoke.py new file mode 100644 index 0000000..553026f --- /dev/null +++ b/tests/e2e/test_minimal_training_smoke.py @@ -0,0 +1,102 @@ +"""End-to-end smoke tests that run tiny training experiments.""" + +from __future__ import annotations + +import os +import subprocess +import sys +from pathlib import Path + +import pytest + +pytestmark = [pytest.mark.e2e, pytest.mark.integration] + + +def _build_pythonpath(root: Path) -> str: + parts = [ + str(root), + str(root / "examples"), + str(root / "third_party"), + str(root / "third_party" / "rw_rl"), + str(root / "third_party" / "rw_rl" / "src"), + ] + existing = os.environ.get("PYTHONPATH", "") + return ":".join(parts + ([existing] if existing else [])) + + +def _run_subprocess(cmd: list[str], root: Path, timeout: int) -> subprocess.CompletedProcess[str]: + env = os.environ.copy() + env["PYTHONPATH"] = _build_pythonpath(root) + env.setdefault("HYDRA_FULL_ERROR", "1") + env.setdefault("PYTHONHASHSEED", "0") + return subprocess.run( + cmd, + cwd=str(root), + env=env, + check=True, + timeout=timeout, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + + +@pytest.mark.parametrize( + ("case_name", "overrides", "timeout_s"), + [ + ("syncrl_local", [], 180), + ( + "syncrl_local_replay", + [ + "buffer.type=ReplayBuffer", + "buffer.sampler.type=uniform", + "train.batch_size=4", + "train.max_rollout_steps=6", + "env.max_episode_steps=6", + ], + 180, + ), + ( + "syncrl_local_multi_worker", + [ + "env.num_workers=2", + "train.batch_size=8", + "train.max_rollout_steps=4", + ], + 180, + ), + ], +) +def test_minimal_training_experiment_smoke(tmp_path, case_name: str, overrides: list[str], timeout_s: int): + root = Path(__file__).resolve().parents[2] + script = root / "tests" / "tests_utils" / "run_mini_experiment.py" + run_dir = tmp_path / case_name + + cmd = [ + sys.executable, + str(script), + "--config-name", + "syncrl_local", + "+train.seed=0", + f"hydra.run.dir={run_dir}", + f"log.log_dir={tmp_path / 'runs'}", + f"log.name={case_name}", + *overrides, + ] + + try: + result = _run_subprocess(cmd, root=root, timeout=timeout_s) + except subprocess.TimeoutExpired: + pytest.fail(f"Minimal experiment {case_name} timed out after {timeout_s}s") + except subprocess.CalledProcessError as exc: + out_tail = (exc.stdout or "").splitlines()[-80:] + err_tail = (exc.stderr or "").splitlines()[-80:] + pytest.fail( + f"Minimal experiment {case_name} failed.\n" + f"Last stdout lines:\n{chr(10).join(out_tail)}\n\n" + f"Last stderr lines:\n{chr(10).join(err_tail)}" + ) + + checkpoint = run_dir / "checkpoints" / "epoch_last.pt" + assert checkpoint.exists(), f"Expected checkpoint at {checkpoint}" + assert "Done." in result.stdout or "Done." in result.stderr diff --git a/tests/e2e/test_openpi_ppo_cli_e2e.py b/tests/e2e/test_openpi_ppo_cli_e2e.py new file mode 100644 index 0000000..72745df --- /dev/null +++ b/tests/e2e/test_openpi_ppo_cli_e2e.py @@ -0,0 +1,134 @@ +"""Minimal end-to-end training tests for OpenPI PPO examples.""" + +from __future__ import annotations + +import subprocess +from pathlib import Path + +import pytest + +from tests.e2e._cli_smoke_utils import ( + artifact_mtime, + assert_subprocess_success, + assert_updated_artifact, + example_python, + resolve_gpu_ids, + run_example, + short_ray_tmpdir, +) + +pytestmark = [pytest.mark.e2e, pytest.mark.integration, pytest.mark.slow] + + +def _assert_openpi_assets(root: Path) -> None: + model_path = Path("/data/ckpts/RLinf/RLinf-Pi0-LIBERO-Spatial-Object-Goal-SFT") + if not model_path.exists(): + pytest.skip(f"OpenPI model path not found: {model_path}") + + libero_path = root / "examples" / "openpi_ppo" / ".venv" / "LIBERO" + if not libero_path.exists(): + pytest.skip(f"LIBERO assets not found: {libero_path}") + + +def _run_openpi_smoke( + *, + config_name: str, + log_name: str, + required_gpus: int, + ray_key: str, + extra_overrides: list[str] | None = None, + timeout_s: int = 2400, +) -> None: + root = Path(__file__).resolve().parents[2] + _assert_openpi_assets(root) + + example_py = example_python(root, "openpi_ppo") + visible_gpus = ",".join(resolve_gpu_ids(required_count=required_gpus, max_used_mb=1024)) + expected_ckpt = root / "runs" / "openpi_ppo" / log_name / "weights" / "epoch_last.pt" + previous_ckpt_mtime = artifact_mtime(expected_ckpt) + + cmd = [ + str(example_py), + "-m", + "examples.openpi_ppo.train_ppo", + "--config-name", + config_name, + "log=tensorboard", + "+debug=False", + "+cluster.ray_address=local", + "+train.seed=0", + "train.max_epochs=1", + "train.max_rollout_steps=2", + "train.warm_up_rollout_steps=2", + "train.batch_size=16", + "train.mini_batch_size=8", + "train.micro_batch_size=8", + "train.update_epoch=1", + "train.rollout_epoch=1", + f"log.name={log_name}", + *(extra_overrides or []), + ] + + try: + result = run_example( + root, + cmd, + timeout_s=timeout_s, + extra_env={ + "CUDA_VISIBLE_DEVICES": visible_gpus, + "RAY_TMPDIR": short_ray_tmpdir(ray_key), + }, + ) + except subprocess.TimeoutExpired: + pytest.fail(f"{log_name} timed out") + + assert_subprocess_success( + result, + log_name, + expected_paths=[expected_ckpt], + ) + assert_updated_artifact(expected_ckpt, previous_ckpt_mtime, log_name) + + +@pytest.mark.gpu +def test_openpi_ppo_tiny_sync_smoke(): + _run_openpi_smoke( + config_name="train_ppo_tiny", + log_name="openpi_tiny_sync_smoke", + required_gpus=3, + ray_key="pi", + extra_overrides=[ + "env.0.num_envs=8", + ], + ) + + +@pytest.mark.gpu +def test_openpi_ppo_sync_smoke(): + _run_openpi_smoke( + config_name="train_ppo", + log_name="openpi_sync_smoke", + required_gpus=3, + ray_key="ps", + extra_overrides=[ + "env.0.num_envs=8", + ], + ) + + +@pytest.mark.gpu +def test_openpi_ppo_tiny_ddp_smoke(): + _run_openpi_smoke( + config_name="train_ppo_tiny_ddp", + log_name="openpi_tiny_ddp_smoke", + required_gpus=8, + ray_key="pd", + extra_overrides=[ + "cluster=4t4e", + "env=libero_x4", + "env.0.num_envs=8", + "train.mini_batch_size=16", + "train.micro_batch_size=4", + ], + timeout_s=3600, + ) diff --git a/tests/e2e/test_openvla_ppo_cli_e2e.py b/tests/e2e/test_openvla_ppo_cli_e2e.py new file mode 100644 index 0000000..cfbcb41 --- /dev/null +++ b/tests/e2e/test_openvla_ppo_cli_e2e.py @@ -0,0 +1,134 @@ +"""OpenVLA PPO CLI smoke tests for minimal full-flow training.""" + +from __future__ import annotations + +import subprocess +from pathlib import Path + +import pytest +from tests.e2e._cli_smoke_utils import ( + artifact_mtime, + assert_subprocess_success, + assert_updated_artifact, + example_python, + require_model_path, + resolve_gpu_ids, + run_example, + short_ray_tmpdir, +) + +pytestmark = [pytest.mark.e2e, pytest.mark.integration, pytest.mark.slow] + +_OPENVLA_DETERMINISTIC_SMOKE_OVERRIDES = [ + # Keep the real end-to-end path, but reduce rollout/reset randomness so + # the smoke cases are less flaky across runs. + "policy.sampling_params.do_sample=False", + "policy.sampling_params.temperature_train=1.0", + "env.0.use_fixed_reset_state_ids=True", +] + + +def _run_openvla_smoke( + *, + config_name: str, + log_name: str, + required_gpus: int, + ray_key: str, + train_micro_batch_size: int, + eval_rollout_override: str = "train.max_eval_rollout_steps=2", + extra_overrides: list[str] | None = None, + timeout_s: int = 3600, +) -> None: + root = Path(__file__).resolve().parents[2] + require_model_path(root / "examples" / "openvla_ppo" / "conf" / "policy" / "openvla_ppo.yaml") + + visible_gpus = ",".join(resolve_gpu_ids(required_count=required_gpus, max_used_mb=1024)) + example_py = example_python(root, "openvla_ppo") + expected_ckpt = root / "runs" / "openvla_ppo" / log_name / "weights" / "epoch_last.pt" + previous_ckpt_mtime = artifact_mtime(expected_ckpt) + + cmd = [ + str(example_py), + "-m", + "examples.openvla_ppo.train_ppo", + "--config-name", + config_name, + "log=tensorboard", + "+debug=False", + "+cluster.ray_address=local", + "+train.seed=0", + "train.max_epochs=1", + "train.max_rollout_steps=2", + eval_rollout_override, + "train.batch_size=16", + "train.mini_batch_size=8", + f"train.micro_batch_size={train_micro_batch_size}", + f"log.name={log_name}", + *(extra_overrides or []), + ] + + try: + result = run_example( + root, + cmd, + timeout_s=timeout_s, + extra_env={ + "CUDA_VISIBLE_DEVICES": visible_gpus, + "RAY_TMPDIR": short_ray_tmpdir(ray_key), + }, + ) + except subprocess.TimeoutExpired: + pytest.fail(f"{log_name} timed out") + + assert_subprocess_success( + result, + log_name, + expected_paths=[expected_ckpt], + ) + assert_updated_artifact(expected_ckpt, previous_ckpt_mtime, log_name) + + +@pytest.mark.gpu +def test_openvla_ppo_ddp_smoke(): + _run_openvla_smoke( + config_name="train_ppo", + log_name="openvla_ddp_smoke", + required_gpus=4, + ray_key="od", + train_micro_batch_size=1, + extra_overrides=[ + "cluster.train_worker_num=2", + "cluster.eval_worker_num=1", + "env.0.num_envs=16", + *_OPENVLA_DETERMINISTIC_SMOKE_OVERRIDES, + ], + ) + + +@pytest.mark.gpu +def test_openvla_ppo_sync_smoke(): + _run_openvla_smoke( + config_name="train_ppo", + log_name="openvla_sync_smoke", + required_gpus=3, + ray_key="os", + train_micro_batch_size=8, + extra_overrides=[ + "env.0.num_envs=16", + *_OPENVLA_DETERMINISTIC_SMOKE_OVERRIDES, + ], + ) + + +@pytest.mark.gpu +def test_openvla_ppo_colocate_x8_smoke(): + _run_openvla_smoke( + config_name="train_ppo_colocate_ddp_x8", + log_name="openvla_colocate_x8_smoke", + required_gpus=8, + ray_key="ox", + train_micro_batch_size=1, + eval_rollout_override="+train.max_eval_rollout_steps=2", + extra_overrides=[*_OPENVLA_DETERMINISTIC_SMOKE_OVERRIDES], + timeout_s=5400, + ) diff --git a/tests/e2e/test_wbc_tracking_cli_e2e.py b/tests/e2e/test_wbc_tracking_cli_e2e.py new file mode 100644 index 0000000..c74be4f --- /dev/null +++ b/tests/e2e/test_wbc_tracking_cli_e2e.py @@ -0,0 +1,206 @@ +"""WBC tracking CLI smoke tests for minimal full-flow training.""" + +from __future__ import annotations + +import subprocess +from pathlib import Path + +import pytest + +from tests.e2e._cli_smoke_utils import ( + assert_subprocess_success, + assert_updated_artifact, + example_python, + find_new_hydra_run_dir, + resolve_gpu_ids, + run_example, + short_ray_tmpdir, + snapshot_hydra_runs, +) + +pytestmark = [pytest.mark.e2e, pytest.mark.integration, pytest.mark.slow] + + +def _wbc_minimal_stability_overrides() -> list[str]: + return [ + "+env.env_kwargs.env_cfg.override.scene.terrain.visual_material=null", + "+env.env_kwargs.env_cfg.override.scene.terrain.physics_material=null", + "+env.env_kwargs.env_cfg.override.scene.contact_forces.debug_vis=false", + "+env.env_kwargs.env_cfg.override.commands.motion.debug_vis=false", + ] + + +def _wbc_multi_task_stability_overrides(env_indices: list[int]) -> list[str]: + overrides: list[str] = [] + for env_idx in env_indices: + prefix = f"+env.{env_idx}.env_kwargs.env_cfg.override" + overrides.extend( + [ + f"{prefix}.scene.terrain.visual_material=null", + f"{prefix}.scene.terrain.physics_material=null", + f"{prefix}.scene.contact_forces.debug_vis=false", + f"{prefix}.commands.motion.debug_vis=false", + ] + ) + return overrides + + +@pytest.mark.gpu +@pytest.mark.isaaclab +def test_wbc_tracking_launch_smoke(): + motion_dir = Path(".data/lafan1/retargeted/wbc_tracking") + if not motion_dir.exists(): + pytest.skip(f"WBC motion asset directory not found: {motion_dir}") + + root = Path(__file__).resolve().parents[2] + example_py = example_python(root, "wbc_tracking") + visible_gpus = ",".join(resolve_gpu_ids(required_count=3, max_used_mb=1024)) + log_name = "wbc_tracking_launch_smoke" + hydra_runs_before = snapshot_hydra_runs(root) + cmd = [ + str(example_py), + str(root / "examples" / "wbc_tracking" / "train.py"), + "--config-name", + "launch", + "debug=False", + "cluster.ray_address=local", + "+train.seed=0", + "env.num_envs=16", + "train.max_rollout_steps=2", + "train.max_epochs=1", + "train.batch_size=16", + f"log.name={log_name}", + *_wbc_minimal_stability_overrides(), + ] + + try: + result = run_example( + root, + cmd, + timeout_s=2400, + extra_env={ + "CUDA_VISIBLE_DEVICES": visible_gpus, + "RAY_TMPDIR": short_ray_tmpdir("w"), + }, + ) + except subprocess.TimeoutExpired: + pytest.fail("wbc_tracking launch smoke test timed out") + + assert_subprocess_success( + result, + "wbc_tracking launch smoke test", + allow_known_wbc_cleanup_error=True, + ) + run_dir = find_new_hydra_run_dir(root, hydra_runs_before, log_name) + checkpoint = run_dir / "checkpoints" / "epoch_last.pt" + assert_updated_artifact( + checkpoint, + previous_mtime=None, + name="wbc_tracking launch smoke test", + ) + + +@pytest.mark.gpu +@pytest.mark.isaaclab +def test_wbc_tracking_ddp_smoke(): + motion_dir = Path(".data/lafan1/retargeted/wbc_tracking") + if not motion_dir.exists(): + pytest.skip(f"WBC motion asset directory not found: {motion_dir}") + + root = Path(__file__).resolve().parents[2] + example_py = example_python(root, "wbc_tracking") + visible_gpus = ",".join(resolve_gpu_ids(required_count=6, max_used_mb=1024)) + log_name = "wbc_tracking_ddp_smoke" + hydra_runs_before = snapshot_hydra_runs(root) + cmd = [ + str(example_py), + str(root / "examples" / "wbc_tracking" / "train.py"), + "--config-name", + "launch_ddp", + "debug=False", + "cluster.ray_address=local", + "+train.seed=0", + "env.0.num_envs=16", + "env.1.num_envs=16", + "train.max_rollout_steps=2", + "train.max_epochs=1", + "train.batch_size=16", + f"log.name={log_name}", + *_wbc_multi_task_stability_overrides([0, 1]), + ] + + try: + result = run_example( + root, + cmd, + timeout_s=3000, + extra_env={ + "CUDA_VISIBLE_DEVICES": visible_gpus, + "RAY_TMPDIR": short_ray_tmpdir("wd"), + }, + ) + except subprocess.TimeoutExpired: + pytest.fail("wbc_tracking ddp smoke test timed out") + + assert_subprocess_success( + result, + "wbc_tracking ddp smoke test", + allow_known_wbc_cleanup_error=True, + ) + run_dir = find_new_hydra_run_dir(root, hydra_runs_before, log_name) + checkpoint = run_dir / "checkpoints" / "epoch_last.pt" + assert_updated_artifact( + checkpoint, + previous_mtime=None, + name="wbc_tracking ddp smoke test", + ) + + +@pytest.mark.gpu +@pytest.mark.isaaclab +def test_wbc_tracking_local_smoke(): + motion_dir = Path(".data/lafan1/retargeted/wbc_tracking") + if not motion_dir.exists(): + pytest.skip(f"WBC motion asset directory not found: {motion_dir}") + + root = Path(__file__).resolve().parents[2] + log_name = "wbc_tracking_local_smoke" + hydra_runs_before = snapshot_hydra_runs(root) + cmd = [ + "bash", + str(root / "examples" / "wbc_tracking" / "launch_local.sh"), + "debug=False", + "env=single_task", + "+train.seed=0", + "env.num_envs=16", + "train.max_epochs=1", + "train.max_rollout_steps=2", + "train.batch_size=16", + f"log.name={log_name}", + *_wbc_minimal_stability_overrides(), + ] + + try: + result = run_example( + root, + cmd, + timeout_s=1800, + extra_env={ + "CUDA_VISIBLE_DEVICES": resolve_gpu_ids(required_count=1, max_used_mb=1024)[0], + "RAY_TMPDIR": short_ray_tmpdir("l"), + }, + ) + except subprocess.TimeoutExpired: + pytest.fail("wbc_tracking local smoke test timed out") + + assert_subprocess_success( + result, + "wbc_tracking local smoke test", + ) + run_dir = find_new_hydra_run_dir(root, hydra_runs_before, log_name) + checkpoint = run_dir / "checkpoints" / "epoch_last.pt" + assert_updated_artifact( + checkpoint, + previous_mtime=None, + name="wbc_tracking local smoke test", + ) diff --git a/tests/integration/__init__.py b/tests/integration/__init__.py new file mode 100644 index 0000000..c210fac --- /dev/null +++ b/tests/integration/__init__.py @@ -0,0 +1 @@ +"""Integration tests.""" diff --git a/tests/integration/buffer/__init__.py b/tests/integration/buffer/__init__.py new file mode 100644 index 0000000..09794ed --- /dev/null +++ b/tests/integration/buffer/__init__.py @@ -0,0 +1 @@ +"""Buffer integration tests.""" diff --git a/tests/integration/buffer/test_buffer_episode_unit.py b/tests/integration/buffer/test_buffer_episode_unit.py new file mode 100644 index 0000000..8238b53 --- /dev/null +++ b/tests/integration/buffer/test_buffer_episode_unit.py @@ -0,0 +1,109 @@ +""" +Integration tests for episode-unit storage in ReplayBuffer. +under episode-unit mode, a full processing flow of storing +and sampling of replaybuffer. +""" + +import gymnasium as gym +import pytest +import torch + +from rlightning.buffer.replay_buffer import ReplayBuffer +from rlightning.types import BatchedData, EnvRet, PolicyResponse +from rlightning.types.metadata import EnvMeta +from rlightning.utils.config import BufferConfig + +from tensordict import TensorDict + +pytestmark = pytest.mark.integration + + +def _make_env_meta(env_id: str) -> EnvMeta: + return EnvMeta( + env_id=env_id, + action_space=gym.spaces.Discrete(2), + observation_space=gym.spaces.Box(low=0, high=1, shape=(1,), dtype=float), + num_envs=1, + ) + + +def _make_buffer_config(): + return BufferConfig.from_dict( + { + "type": "ReplayBuffer", + "capacity": 10, + "sampler": {"type": "uniform"}, + "storage": {"type": "unified", "device": "cpu", "unit": "episode"}, + } + ) + + +def test_episode_unit_storage_flow(): + tensordict = pytest.importorskip("tensordict") + + env_id = "env-episode" + buffer = ReplayBuffer(config=_make_buffer_config()) + buffer.init([_make_env_meta(env_id)], [env_id]) + + # Episode 1: 3 steps -> stored episode length 2 after postprocess + for step in range(3): + env_ret = EnvRet( + env_id=env_id, + observation=float(step), + last_reward=float(step), + last_terminated=step == 2, + last_truncated=False, + info={}, + ) + policy_resp = PolicyResponse(env_id=env_id, action=float(step)) + buffer.add_batched_transition( + BatchedData([env_id], [env_ret]), + BatchedData([env_id], [policy_resp]), + ) + buffer.truncate_episodes([env_id]) + + # Episode 2: same length + for step in range(3): + env_ret = EnvRet( + env_id=env_id, + observation=float(step + 10), + last_reward=float(step), + last_terminated=step == 2, + last_truncated=False, + info={}, + ) + policy_resp = PolicyResponse(env_id=env_id, action=float(step)) + buffer.add_batched_transition( + BatchedData([env_id], [env_ret]), + BatchedData([env_id], [policy_resp]), + ) + buffer.truncate_episodes([env_id]) + + assert len(buffer) == 2 + + sample_data = buffer.sample(batch_size=2) + assert isinstance(sample_data[0], dict) + sample_data = [TensorDict.from_dict(data, auto_batch_size=True) for data in sample_data] + assert sample_data[0].batch_size == torch.Size([2, 2]) + + # Content validation: verify sampled data matches written episodes + sampled_obs = sample_data[0]["observation"] # shape: [2, 2] (batch=2, time=2) + sampled_action = sample_data[0]["action"] # shape: [2, 2] + + # Episode 1: obs=[0.0, 1.0], action=[0.0, 1.0] (after postprocess, length 3->2) + # Episode 2: obs=[10.0, 11.0], action=[0.0, 1.0] + # Each sampled episode should match one of these patterns + valid_first_obs = {0.0, 10.0} + for i in range(2): + first_obs = sampled_obs[i, 0].item() + assert first_obs in valid_first_obs, f"Unexpected first obs: {first_obs}" + + # Verify temporal consistency within episode + if first_obs == 0.0: + assert sampled_obs[i, 1].item() == 1.0, "Episode 1 obs sequence mismatch" + assert sampled_action[i, 0].item() == 0.0, "Episode 1 action mismatch" + assert sampled_action[i, 1].item() == 1.0, "Episode 1 action mismatch" + else: + assert sampled_obs[i, 1].item() == 11.0, "Episode 2 obs sequence mismatch" + assert sampled_action[i, 0].item() == 0.0, "Episode 2 action mismatch" + assert sampled_action[i, 1].item() == 1.0, "Episode 2 action mismatch" diff --git a/tests/integration/buffer/test_buffer_episode_unit_variable_length.py b/tests/integration/buffer/test_buffer_episode_unit_variable_length.py new file mode 100644 index 0000000..1988c12 --- /dev/null +++ b/tests/integration/buffer/test_buffer_episode_unit_variable_length.py @@ -0,0 +1,76 @@ +""" +Integration test for episode-unit storage with variable episode lengths. +""" + +import gymnasium as gym +import pytest + +from rlightning.buffer.replay_buffer import ReplayBuffer +from rlightning.types import BatchedData, EnvRet, PolicyResponse +from rlightning.types.metadata import EnvMeta +from rlightning.utils.config import BufferConfig + +pytestmark = pytest.mark.integration + + +def _make_env_meta(env_id: str) -> EnvMeta: + return EnvMeta( + env_id=env_id, + action_space=gym.spaces.Discrete(2), + observation_space=gym.spaces.Box(low=0, high=1, shape=(1,), dtype=float), + num_envs=1, + ) + + +def _make_buffer_config(): + return BufferConfig.from_dict( + { + "type": "ReplayBuffer", + "capacity": 20, + "sampler": {"type": "all"}, + "storage": {"type": "unified", "device": "cpu", "unit": "episode"}, + } + ) + + +def _add_episode(buffer, env_id: str, steps: int): + for step in range(steps): + env_ret = EnvRet( + env_id=env_id, + observation=float(step), + last_reward=float(step), + last_terminated=step == (steps - 1), + last_truncated=False, + info={}, + ) + policy_resp = PolicyResponse(env_id=env_id, action=float(step)) + buffer.add_batched_transition( + BatchedData([env_id], [env_ret]), + BatchedData([env_id], [policy_resp]), + ) + buffer.truncate_episodes([env_id]) + + +def test_episode_unit_variable_length_returns_list(): + tensordict = pytest.importorskip("tensordict") + + env_id = "env-episode-varlen" + buffer = ReplayBuffer(config=_make_buffer_config()) + buffer.init([_make_env_meta(env_id)], [env_id]) + + # Episode lengths 3 and 4 steps -> stored transition lengths 2 and 3 + _add_episode(buffer, env_id, steps=3) + _add_episode(buffer, env_id, steps=4) + + assert len(buffer) == 2 + + sample_data = buffer.sample(batch_size=None) + assert len(sample_data) == 1 + + # Variable lengths should produce a list of TensorDicts + assert isinstance(sample_data[0], list) + assert len(sample_data[0]) == 2 + assert all(isinstance(item, tensordict.TensorDict) for item in sample_data[0]) + + lengths = [len(item) for item in sample_data[0]] + assert sorted(lengths) == [2, 3] diff --git a/tests/integration/buffer/test_buffer_integration_flow.py b/tests/integration/buffer/test_buffer_integration_flow.py new file mode 100644 index 0000000..a2bb722 --- /dev/null +++ b/tests/integration/buffer/test_buffer_integration_flow.py @@ -0,0 +1,132 @@ +""" +Integration tests for buffer components working together. +""" + +import gymnasium as gym +import pytest + +from rlightning.buffer.replay_buffer import ReplayBuffer +from rlightning.types import BatchedData, EnvRet, PolicyResponse +from rlightning.types.metadata import EnvMeta +from rlightning.utils.config import BufferConfig + +pytestmark = pytest.mark.integration + + +def _make_buffer_config(auto_truncate=False): + return BufferConfig.from_dict( + { + "type": "ReplayBuffer", + "capacity": 8, + "auto_truncate_episode": auto_truncate, + "sampler": {"type": "uniform"}, + "storage": {"type": "unified", "device": "cpu", "unit": "transition"}, + } + ) + + +def _make_env_meta(env_id: str) -> EnvMeta: + return EnvMeta( + env_id=env_id, + action_space=gym.spaces.Discrete(2), + observation_space=gym.spaces.Box(low=0, high=1, shape=(1,), dtype=float), + num_envs=1, + ) + + +def test_replay_buffer_transition_flow(): + env_id = "env-1" + buffer = ReplayBuffer(config=_make_buffer_config()) + buffer.init([_make_env_meta(env_id)], [env_id]) + + for step in range(2): + env_ret = EnvRet( + env_id=env_id, + observation=float(step), + last_reward=float(step), + last_terminated=step == 1, + last_truncated=False, + info={}, + ) + policy_resp = PolicyResponse(env_id=env_id, action=float(step)) + buffer.add_batched_transition( + BatchedData([env_id], [env_ret]), + BatchedData([env_id], [policy_resp]), + ) + + buffer.truncate_episodes([env_id]) + + # Precise length check: 2 steps -> 1 transition after postprocess + assert len(buffer) == 1 + + sample_data = buffer.sample(batch_size=1) + assert len(sample_data) == 1 + assert isinstance(sample_data[0], dict) + + # Keys validation + keys = set(sample_data[0].keys()) + assert "observation" in keys + assert "next_observation" in keys + assert "action" in keys + assert "reward" in keys + + # Content validation: verify sampled data matches written values + sample = sample_data[0] + # Step 0: obs=0.0, action=0.0 -> Step 1: obs=1.0, reward=1.0 + assert sample["observation"].item() == 0.0, "observation should be from step 0" + assert sample["next_observation"].item() == 1.0, "next_observation should be from step 1" + assert sample["action"].item() == 0.0, "action should be from step 0" + assert sample["reward"].item() == 1.0, "reward should be from step 1" + + +def test_replay_buffer_async_data_flow(): + env_id = "env-async" + buffer = ReplayBuffer(config=_make_buffer_config(auto_truncate=True)) + buffer.init([_make_env_meta(env_id)], [env_id]) + + env_ret_0 = EnvRet( + env_id=env_id, + observation=0.0, + last_reward=1.0, + last_terminated=False, + last_truncated=False, + info={}, + ) + policy_resp_0 = PolicyResponse(env_id=env_id, action=0.0) + + env_ret_1 = EnvRet( + env_id=env_id, + observation=1.0, + last_reward=2.0, + last_terminated=True, + last_truncated=False, + info={}, + ) + policy_resp_1 = PolicyResponse(env_id=env_id, action=1.0) + + buffer.add_batched_data_async(BatchedData([env_id], [env_ret_0])) + buffer.add_batched_data_async(BatchedData([env_id], [policy_resp_0])) + buffer.add_batched_data_async(BatchedData([env_id], [env_ret_1])) + buffer.add_batched_data_async(BatchedData([env_id], [policy_resp_1])) + + # Precise length check: 2 steps -> 1 transition after postprocess + assert len(buffer) == 1 + + sample_data = buffer.sample(batch_size=1) + assert isinstance(sample_data[0], dict) + + # Keys validation (consistent with test_replay_buffer_transition_flow) + keys = set(sample_data[0].keys()) + assert "observation" in keys + assert "next_observation" in keys + assert "action" in keys + assert "reward" in keys + + # Content validation: verify sampled data matches written values + sample = sample_data[0] + # env_ret_0: obs=0.0, reward=1.0, action=0.0 + # env_ret_1: obs=1.0, reward=2.0, action=1.0, terminated=True + assert sample["observation"].item() == 0.0, "observation should be from env_ret_0" + assert sample["next_observation"].item() == 1.0, "next_observation should be from env_ret_1" + assert sample["action"].item() == 0.0, "action should be from policy_resp_0" + assert sample["reward"].item() == 2.0, "reward should be from env_ret_1" diff --git a/tests/integration/buffer/test_buffer_multi_env_and_sampler.py b/tests/integration/buffer/test_buffer_multi_env_and_sampler.py new file mode 100644 index 0000000..dfc4727 --- /dev/null +++ b/tests/integration/buffer/test_buffer_multi_env_and_sampler.py @@ -0,0 +1,161 @@ +""" +Deeper integration tests for buffer sampling and stats. +""" + +import gymnasium as gym +import pytest +from tensordict import TensorDict + +from rlightning.buffer.replay_buffer import ReplayBuffer +from rlightning.types import BatchedData, EnvRet, PolicyResponse +from rlightning.types.metadata import EnvMeta +from rlightning.utils.config import BufferConfig + +pytestmark = pytest.mark.integration + + +def _make_env_meta(env_id: str) -> EnvMeta: + return EnvMeta( + env_id=env_id, + action_space=gym.spaces.Discrete(2), + observation_space=gym.spaces.Box(low=0, high=1, shape=(1,), dtype=float), + num_envs=1, + ) + + +def _make_buffer_config(sampler_type: str = "uniform"): + return BufferConfig.from_dict( + { + "type": "ReplayBuffer", + "capacity": 16, + "sampler": {"type": sampler_type}, + "storage": {"type": "unified", "device": "cpu", "unit": "transition"}, + } + ) + + +def test_replay_buffer_multi_env_sampling(): + env_ids = ["env-1", "env-2"] + buffer = ReplayBuffer(config=_make_buffer_config()) + buffer.init([_make_env_meta(env_id) for env_id in env_ids], env_ids) + + # Two steps for each env (episode length=2 -> 1 transition per episode) + for step in range(2): + env_rets = [] + policy_resps = [] + for env_id in env_ids: + env_rets.append( + EnvRet( + env_id=env_id, + observation=float(step), + last_reward=float(step), + last_terminated=step == 1, + last_truncated=False, + info={}, + ) + ) + policy_resps.append(PolicyResponse(env_id=env_id, action=float(step))) + buffer.add_batched_transition( + BatchedData(env_ids, env_rets), + BatchedData(env_ids, policy_resps), + ) + + buffer.truncate_episodes(env_ids) + + # 2 envs × 1 transition each = 2 transitions total + assert len(buffer) == 2 + sample_data = buffer.sample(batch_size=2) + + assert len(sample_data) == 1 + assert isinstance(sample_data[0], dict) + assert len(TensorDict.from_dict(sample_data[0], auto_batch_size=True)) == 2 + + # Content validation: verify sampled data matches written values + # Both envs wrote obs=0.0 at step 0, next_obs=1.0 at step 1 + for i in range(2): + assert sample_data[0]["observation"][i].item() == 0.0, "observation should be 0.0" + assert sample_data[0]["next_observation"][i].item() == 1.0, "next_observation should be 1.0" + assert sample_data[0]["action"][i].item() == 0.0, "action should be 0.0" + assert sample_data[0]["reward"][i].item() == 1.0, "reward should be 1.0" + + +def test_replay_buffer_all_sampler_returns_all_data(): + env_id = "env-all" + buffer = ReplayBuffer(config=_make_buffer_config(sampler_type="all")) + buffer.init([_make_env_meta(env_id)], [env_id]) + + for step in range(2): + env_ret = EnvRet( + env_id=env_id, + observation=float(step), + last_reward=float(step), + last_terminated=step == 1, + last_truncated=False, + info={}, + ) + policy_resp = PolicyResponse(env_id=env_id, action=float(step)) + buffer.add_batched_transition( + BatchedData([env_id], [env_ret]), + BatchedData([env_id], [policy_resp]), + ) + + buffer.truncate_episodes([env_id]) + + # all sampler ignores batch_size and returns everything + sample_data = buffer.sample(batch_size=None) + + assert len(sample_data) == 1 + assert isinstance(sample_data[0], dict) + assert len(TensorDict.from_dict(sample_data[0], auto_batch_size=True)) == len(buffer) + + # Content validation: verify sampled data matches written values + # 2 steps -> 1 transition: obs=0.0, next_obs=1.0, action=0.0, reward=1.0 + assert sample_data[0]["observation"].item() == 0.0, "observation should be 0.0" + assert sample_data[0]["next_observation"].item() == 1.0, "next_observation should be 1.0" + assert sample_data[0]["action"].item() == 0.0, "action should be 0.0" + assert sample_data[0]["reward"].item() == 1.0, "reward should be 1.0" + + +def test_replay_buffer_batch_sampler_no_replacement(): + tensordict = pytest.importorskip("tensordict") + + env_id = "env-batch" + buffer = ReplayBuffer(config=_make_buffer_config(sampler_type="batch")) + buffer.init([_make_env_meta(env_id)], [env_id]) + + # Episode length=5 -> 4 transitions stored after truncate. + for step in range(5): + env_ret = EnvRet( + env_id=env_id, + observation=float(step), + last_reward=float(step), + last_terminated=step == 4, + last_truncated=False, + info={}, + ) + policy_resp = PolicyResponse(env_id=env_id, action=float(step)) + buffer.add_batched_transition( + BatchedData([env_id], [env_ret]), + BatchedData([env_id], [policy_resp]), + ) + + buffer.truncate_episodes([env_id]) + + # 5 steps -> 4 transitions after postprocess + assert len(buffer) == 4 + + sample_data = buffer.sample(batch_size=3) + assert isinstance(sample_data[0], dict) + + # Batch sampler should return unique samples (no replacement) + sample_obs = sample_data[0]["observation"] + assert len(set(sample_obs)) == len(sample_obs), "batch sampler should not have duplicates" + + # Content validation: sampled obs should be from written values [0.0, 1.0, 2.0, 3.0] + valid_obs = {0.0, 1.0, 2.0, 3.0} + for obs in sample_obs: + assert obs.item() in valid_obs, f"unexpected observation {obs}" + + # Requesting more samples than buffer size should raise ValueError + with pytest.raises(ValueError): + buffer.sample(batch_size=5) diff --git a/tests/integration/buffer/test_buffer_multi_worker_sampling.py b/tests/integration/buffer/test_buffer_multi_worker_sampling.py new file mode 100644 index 0000000..739098e --- /dev/null +++ b/tests/integration/buffer/test_buffer_multi_worker_sampling.py @@ -0,0 +1,94 @@ +""" +Integration test for buffer sampling across multiple train workers. +""" + +import gymnasium as gym +import pytest + +from rlightning.buffer.replay_buffer import ReplayBuffer +from rlightning.types import BatchedData, EnvRet, PolicyResponse +from rlightning.types.metadata import EnvMeta +from rlightning.utils.config import BufferConfig + +pytestmark = pytest.mark.integration + + +def _make_env_meta(env_id: str) -> EnvMeta: + return EnvMeta( + env_id=env_id, + action_space=gym.spaces.Discrete(2), + observation_space=gym.spaces.Box(low=0, high=1, shape=(1,), dtype=float), + num_envs=1, + ) + + +def _make_buffer_config(): + return BufferConfig.from_dict( + { + "type": "ReplayBuffer", + "capacity": 16, + "sampler": {"type": "uniform"}, + "storage": {"type": "unified", "device": "cpu", "unit": "transition"}, + } + ) + + +def test_multi_worker_sampling_split(): + tensordict = pytest.importorskip("tensordict") + + env_id = "env-multi-worker" + buffer = ReplayBuffer(config=_make_buffer_config()) + buffer.init([_make_env_meta(env_id)], [env_id]) + + # Two episodes with 4 transitions each -> total 8 transitions + for episode_idx in range(2): + for step in range(5): + env_ret = EnvRet( + env_id=env_id, + observation=float(step + episode_idx * 10), + last_reward=float(step), + last_terminated=step == 4, + last_truncated=False, + info={}, + ) + policy_resp = PolicyResponse(env_id=env_id, action=float(step)) + buffer.add_batched_transition( + BatchedData([env_id], [env_ret]), + BatchedData([env_id], [policy_resp]), + ) + buffer.truncate_episodes([env_id]) + + assert len(buffer) == 8 + + # Simulate two train workers bound to the same storage. + buffer.table._storage_to_train_workers = {0: [0, 1]} + + sample_data = buffer.sample(batch_size=6, shuffle=False, drop_last=True) + + assert len(sample_data) == 2 + assert isinstance(sample_data[0], dict) + assert isinstance(sample_data[1], dict) + + # Each worker should receive half the samples (6 total -> 3 each). + assert len(next(iter(sample_data[0].values()))) == 3 + assert len(next(iter(sample_data[1].values()))) == 3 + + # Ensure workers do not receive identical data sets. + obs_worker0 = sample_data[0]["observation"].tolist() + obs_worker1 = sample_data[1]["observation"].tolist() + assert obs_worker0 != obs_worker1 + + # Content validation: all sampled obs should come from written data + # Episode 0: obs in [0,1,2,3], Episode 1: obs in [10,11,12,13] + valid_obs = {0.0, 1.0, 2.0, 3.0, 10.0, 11.0, 12.0, 13.0} + all_sampled_obs = obs_worker0 + obs_worker1 + for obs in all_sampled_obs: + assert obs in valid_obs, f"Unexpected observation {obs}" + + # Verify next_observation corresponds to obs + 1 (within same episode) + for worker_data in sample_data: + obs_list = worker_data["observation"].tolist() + next_obs_list = worker_data["next_observation"].tolist() + for obs, next_obs in zip(obs_list, next_obs_list): + # obs and next_obs should differ by 1 (same episode continuity) + assert next_obs == obs + 1.0, f"next_obs {next_obs} should be obs {obs} + 1" diff --git a/tests/integration/buffer/test_replay_buffer_with_env.py b/tests/integration/buffer/test_replay_buffer_with_env.py new file mode 100644 index 0000000..f3d14fb --- /dev/null +++ b/tests/integration/buffer/test_replay_buffer_with_env.py @@ -0,0 +1,416 @@ +import os +import random + +import pytest +import torch + +from rlightning.buffer.utils.preprocessors import get_preprocessor_cls +from rlightning.env import EnvGroup +from rlightning.types import BatchedData, EnvMeta, PolicyResponse +from rlightning.utils.builders import build_data_buffer, build_env_group +from rlightning.utils.config import BufferConfig +from rlightning.utils.ray import TaskSubmitter, resolve_object +from rlightning.utils.utils import InternalFlag +from tests.test_utils import _make_env_config, require_env_backend + +pytestmark = pytest.mark.integration + + +@pytest.fixture(scope="module", autouse=True) +def _auto_ray_cluster(): + # Ensure a Ray cluster is available for this module's tests + from tests.test_utils import setup_ray_cluster, teardown_ray_cluster + + setup_ray_cluster() + yield + teardown_ray_cluster() + + +def _make_buffer_config(capacity: int = 100) -> BufferConfig: + return BufferConfig.from_dict( + { + "type": "ReplayBuffer", + "capacity": capacity, + "storage": { + "type": "unified", + "device": "cpu", + }, + "sampler": { + "type": "uniform", + }, + } + ) + + +def _make_env(env_cfg, as_remote: bool): + require_env_backend(env_cfg.backend) + os.environ["RLIGHTNING_REMOTE_ENV"] = "1" if as_remote else "0" + assert InternalFlag.REMOTE_ENV == as_remote + env, _ = EnvGroup.build_env(env_cfg) + if as_remote: + env_meta = env.init.remote() + else: + env_meta = env.init() + return env, env_meta + + +def _make_env_group(num_env_workers, as_remote): + os.environ["RLIGHTNING_REMOTE_ENV"] = "1" if as_remote else "0" + assert InternalFlag.REMOTE_ENV == as_remote + require_env_backend("mujoco") + env_cfgs = [ + _make_env_config(env_backend="mujoco", task_name="Ant-v5", num_workers=num_env_workers), + ] + + env_group = build_env_group(env_cfgs) + + return env_group + + +def _mock_batched_policy_ret(action_spaces, env_ids): + action_list = [resolve_object(action_space).sample() for action_space in action_spaces] + policy_resp_list = [ + PolicyResponse(env_id=env_id, action=action, log_prob=0.0, entropy=0.0) + for action, env_id in zip(action_list, env_ids) + ] + batched_policy_resp = BatchedData(env_ids, policy_resp_list) + return batched_policy_resp + + +def _mock_policy_ret(action_space, env_id): + action = action_space.sample() + policy_resp = PolicyResponse(env_id=env_id, action=action, log_prob=0.0, entropy=0.0) + return policy_resp + + +def _make_buffer(buffer_cfg, env_meta_list=None, as_remote: bool = False): + os.environ["RLIGHTNING_REMOTE_STORAGE"] = "1" if as_remote else "0" + assert InternalFlag.REMOTE_STORAGE == as_remote + + env_meta = env_meta_list[0] if env_meta_list is not None else None + if env_meta is not None and env_meta.observation_space is not None: + observation_space = env_meta.observation_space + obs_preprocessor = get_preprocessor_cls(observation_space)(observation_space) + buffer = build_data_buffer(buffer_cls=buffer_cfg.type, buffer_cfg=buffer_cfg, obs_preprocessor=obs_preprocessor) + else: + buffer = build_data_buffer(buffer_cls=buffer_cfg.type, buffer_cfg=buffer_cfg) + + buffer.init(env_meta_list) + + return buffer + + +def _make_episode(episode_length: int = 8, n_envs: int = 1, obs_dim: int = 3, act_dim: int = 2): + if n_envs == 1: + return { + "obs": torch.randn(episode_length, obs_dim), + "action": torch.randn(episode_length, act_dim), + } + + return { + "obs": torch.randn(episode_length, n_envs, obs_dim), + "action": torch.randn(episode_length, n_envs, act_dim), + } + + +@pytest.mark.parametrize( + "env_backend,task_name", + [ + ["mujoco", "Ant-v5"], + ["ale", "ALE/Adventure-v5"], + ], +) +@pytest.mark.parametrize("num_env_workers", [1, 2]) +@pytest.mark.parametrize("env_as_remote", [False]) +@pytest.mark.parametrize("buffer_as_remote", [False]) +def test_replay_buffer( + env_backend: str, + task_name: str, + num_env_workers: int, + env_as_remote: bool, + buffer_as_remote: bool, +): + """Test the functionality of replay buffer. + + The coverage includes: + + 1. Making timesteps holder (now only test for single-environment yet) + 2. Functional preprocessing for each timestep + 3. Functional postprocessing for a dict of batched-timesteps, in the order of [N_timesteps, N_envs, *inner_dims] + 4. Pushing a dict of batch, but not test for the case of overflow + 5. Functional and Correctness checking for the buffer sampling, cases including: a) predefined batch size; b) random batch size; c) given indices with random batch_size + + Args: + env_backend (str): Environmnt backend + task_name (str): Task name + num_env_workers (int): Environment Instance number + """ + + env_cfg = _make_env_config(env_backend, task_name, num_env_workers) + buffer_cfg = _make_buffer_config(capacity=100) + # TODO(ming): apply environment vectorization for future work + env, env_meta = _make_env(env_cfg, env_as_remote) + env_meta_list = None if env_meta is None else [env_meta] + buffer = _make_buffer(buffer_cfg, env_meta_list, buffer_as_remote) + + batch_size: int = 10 + + env_ret = env.reset() + buffer.add_data_async(env_ret.env_id, env_ret) + step_cnt = 0 + while not (env_ret.last_terminated or env_ret.last_truncated or step_cnt >= env_cfg.max_episode_steps): + action = env.action_space.sample() + policy_resp = PolicyResponse(env_id=env_ret.env_id, action=action, log_prob=0.0, entropy=0.0) + buffer.add_data_async(env_ret.env_id, policy_resp) + + env_ret = env.step(policy_resp) + buffer.add_data_async(env_ret.env_id, env_ret) + + step_cnt += 1 + + action = env.action_space.sample() + policy_resp = PolicyResponse(env_id=env_ret.env_id, action=action, log_prob=0.0, entropy=0.0) + buffer.add_data_async(env_ret.env_id, policy_resp) + buffer.truncate_episodes([policy_resp.env_id]) + + assert env_ret.last_terminated or env_ret.last_truncated, ( + env_ret.last_terminated, + env_ret.last_truncated, + ) + + assert len(buffer) == step_cnt, (step_cnt, len(buffer)) + + custom_batch_size = random.choice(range(1, batch_size)) + sample_data = buffer.sample(batch_size=custom_batch_size) + assert len(sample_data) == 1 + data = resolve_object(sample_data[0]) + for k, v in data.items(): + assert len(v) == custom_batch_size, (k, v.shape) + + +@pytest.mark.parametrize("buffer_as_remote", [True, False]) +@pytest.mark.parametrize("env_as_remote", [False]) +def test_init_creates_storage(buffer_as_remote: bool, env_as_remote: bool): + env_backend = "mujoco" + task_name = "Ant-v5" + num_env_workers = 1 + env_cfg = _make_env_config(env_backend, task_name, num_env_workers) + buffer_cfg = _make_buffer_config(capacity=100) + + _, env_meta = _make_env(env_cfg, env_as_remote) + env_meta_list = None if env_meta is None else [env_meta] + buffer = _make_buffer(buffer_cfg, env_meta_list, buffer_as_remote) + assert len(buffer.storages) > 0 + # assert buf.table is not None + assert buffer.size() == 0 + + +@pytest.mark.parametrize("num_envs", [1, 2, 4]) +@pytest.mark.parametrize("buffer_as_remote", [True, False]) +def test_add_increases_size(num_envs, buffer_as_remote): + env_meta = EnvMeta( + env_id=None, + action_space=None, + observation_space=None, + num_envs=num_envs, + ) + buffer_cfg = _make_buffer_config(capacity=100) + buf = _make_buffer(buffer_cfg, [env_meta], as_remote=buffer_as_remote) + episode_length = 8 + episode = _make_episode(episode_length=episode_length, n_envs=num_envs) + buf.add_episode(episode, num_envs=num_envs) + + assert buf.size() == episode_length * num_envs + + sample_data = buf.sample(batch_size=8) + assert len(sample_data) == 1 + data = resolve_object(sample_data[0]) + for k, v in data.items(): + assert len(v) == 8, (k, v.shape) + + +@pytest.mark.parametrize("buffer_as_remote", [True, False]) +def test_clear_resets_size(buffer_as_remote): + buffer_cfg = _make_buffer_config(capacity=100) + buf = _make_buffer(buffer_cfg, as_remote=buffer_as_remote) + episode_length = 8 + episode = _make_episode(episode_length=episode_length, n_envs=1) + buf.add_episode(episode, num_envs=1) + + assert buf.size() == episode_length + buf.clear() + assert buf.size() == 0 + + +@pytest.mark.parametrize("buffer_as_remote", [True, False]) +def test_circular_mode_overwrite_when_exceed_capacity(buffer_as_remote): + # capacity 10, add 14 => size should be capped at 10 + buffer_cfg = _make_buffer_config(capacity=10) + buf = _make_buffer(buffer_cfg, as_remote=buffer_as_remote) + buf.add_episode(_make_episode(episode_length=14, n_envs=1), num_envs=1) + + assert buf.size() == 10 + + +@pytest.mark.parametrize("env_as_remote", [True, False]) +@pytest.mark.parametrize("buffer_as_remote", [True, False]) +def test_add_transition(env_as_remote, buffer_as_remote): + env_backend = "mujoco" + task_name = "Ant-v5" + num_env_workers = 1 + env_cfg = _make_env_config(env_backend, task_name, num_env_workers) + task_submitter = TaskSubmitter() + + env, _ = _make_env(env_cfg, env_as_remote) + buffer_cfg = _make_buffer_config(capacity=10) + buffer = _make_buffer(buffer_cfg, as_remote=buffer_as_remote) + + env_ret = task_submitter.submit(env.reset, _block=True) + for _ in range(5): + action = task_submitter.submit(env.get_action_space, _block=True).sample() + policy_resp = PolicyResponse(env_id=env_ret.env_id, action=action, log_prob=0.0, entropy=0.0) + buffer.add_transition(env_ret.env_id, env_ret, policy_resp) + env_ret = task_submitter.submit(env.step, policy_resp, _block=True) + + action = task_submitter.submit(env.get_action_space, _block=True).sample() + policy_resp = PolicyResponse(env_id=env_ret.env_id, action=action, log_prob=0.0, entropy=0.0) + buffer.add_transition(env_ret.env_id, env_ret, policy_resp, truncated=True) + + assert len(buffer) == 5 + + +@pytest.mark.parametrize("env_as_remote", [True, False]) +@pytest.mark.parametrize("buffer_as_remote", [True, False]) +def test_add_batched_transition(env_as_remote, buffer_as_remote): + env_group = _make_env_group(num_env_workers=2, as_remote=env_as_remote) + + buffer_cfg = _make_buffer_config(capacity=100) + buffer = _make_buffer(buffer_cfg, as_remote=buffer_as_remote) + + batched_env_ret, _ = env_group.reset() + for _ in range(5): + batched_policy_resp = _mock_batched_policy_ret(env_group.get_action_spaces(), batched_env_ret.ids()) + buffer.add_batched_transition(batched_env_ret, batched_policy_resp) + batched_env_ret, _ = env_group.step(batched_policy_resp) + + batched_policy_resp = _mock_batched_policy_ret(env_group.get_action_spaces(), batched_env_ret.ids()) + truncations = [True] * len(batched_env_ret) + buffer.add_batched_transition(batched_env_ret, batched_policy_resp, truncations) + + assert len(buffer) == 5 * len(env_group) + + +@pytest.mark.parametrize("env_as_remote", [True, False]) +@pytest.mark.parametrize("buffer_as_remote", [True, False]) +def test_add_data_async(env_as_remote, buffer_as_remote): + env_backend = "mujoco" + task_name = "Ant-v5" + num_env_workers = 1 + env_cfg = _make_env_config(env_backend, task_name, num_env_workers) + + task_submitter = TaskSubmitter() + + env, _ = _make_env(env_cfg, env_as_remote) + buffer_cfg = _make_buffer_config(capacity=10) + buffer = _make_buffer(buffer_cfg, as_remote=buffer_as_remote) + + env_ret = task_submitter.submit(env.reset, _block=True) + buffer.add_data_async(env_ret.env_id, env_ret) + for _ in range(5): + policy_resp = _mock_policy_ret(task_submitter.submit(env.get_action_space, _block=True), env_ret.env_id) + buffer.add_data_async(env_ret.env_id, policy_resp) + env_ret = task_submitter.submit(env.step, policy_resp, _block=True) + buffer.add_data_async(env_ret.env_id, env_ret) + + policy_resp = _mock_policy_ret(task_submitter.submit(env.get_action_space, _block=True), env_ret.env_id) + buffer.add_data_async(env_ret.env_id, policy_resp, truncated=True) + + assert len(buffer) == 5 + + +@pytest.mark.parametrize("env_as_remote", [True]) +@pytest.mark.parametrize("buffer_as_remote", [True, False]) +def test_add_batched_data_async(env_as_remote, buffer_as_remote): + env_group = _make_env_group(num_env_workers=2, as_remote=env_as_remote) + + buffer_cfg = _make_buffer_config(capacity=100) + buffer = _make_buffer(buffer_cfg, as_remote=buffer_as_remote) + + batched_env_ret, _ = env_group.reset() + buffer.add_batched_data_async(batched_env_ret) + for _ in range(5): + batched_policy_resp = _mock_batched_policy_ret(env_group.get_action_spaces(), batched_env_ret.ids()) + env_group.step_async(batched_policy_resp) + buffer.add_batched_data_async(batched_policy_resp) + batched_env_ret, _ = env_group.collect_async() + buffer.add_batched_data_async(batched_env_ret) + + batched_policy_resp = _mock_batched_policy_ret(env_group.get_action_spaces(), batched_env_ret.ids()) + truncations = [True] * len(batched_env_ret) + buffer.add_batched_data_async(batched_policy_resp, truncations) + + # to avoid the case that some of the env's step has not been collected + batche_env_ret, _ = env_group.collect_async(wait_all=True) + if len(batche_env_ret) > 0: + buffer.add_batched_data_async(batche_env_ret) + batched_policy_resp = _mock_batched_policy_ret(env_group.get_action_spaces(), batched_env_ret.ids()) + buffer.add_batched_data_async(batched_policy_resp) + + truncations = [True] * len(batched_env_ret) + buffer.add_batched_data_async(batched_policy_resp, truncations) + + assert len(buffer) == 5 * len(env_group) + + +@pytest.mark.parametrize("env_as_remote", [True, False]) +@pytest.mark.parametrize("buffer_as_remote", [True, False]) +def test_truncate_one_episode(env_as_remote, buffer_as_remote): + env_backend = "mujoco" + task_name = "Ant-v5" + num_env_workers = 1 + env_cfg = _make_env_config(env_backend, task_name, num_env_workers) + + env, _ = _make_env(env_cfg, env_as_remote) + buffer_cfg = _make_buffer_config(capacity=10) + buffer = _make_buffer(buffer_cfg, as_remote=buffer_as_remote) + task_submitter = TaskSubmitter() + + env_ret = task_submitter.submit(env.reset, _block=True) + for _ in range(5): + action = task_submitter.submit(env.get_action_space, _block=True).sample() + policy_resp = PolicyResponse(env_id=env_ret.env_id, action=action, log_prob=0.0, entropy=0.0) + buffer.add_transition(env_ret.env_id, env_ret, policy_resp) + env_ret = task_submitter.submit(env.step, policy_resp, _block=True) + + action = task_submitter.submit(env.get_action_space, _block=True).sample() + policy_resp = PolicyResponse(env_id=env_ret.env_id, action=action, log_prob=0.0, entropy=0.0) + buffer.add_transition(env_ret.env_id, env_ret, policy_resp) + + assert len(buffer) == 0 + + buffer.truncate_one_episode(policy_resp.env_id) + assert len(buffer) == 5 + + +@pytest.mark.parametrize("env_as_remote", [True, False]) +@pytest.mark.parametrize("buffer_as_remote", [True, False]) +def test_truncate_episodes(env_as_remote, buffer_as_remote): + env_group = _make_env_group(num_env_workers=2, as_remote=env_as_remote) + + buffer_cfg = _make_buffer_config(capacity=100) + buffer = _make_buffer(buffer_cfg, as_remote=buffer_as_remote) + + batched_env_ret, _ = env_group.reset() + for _ in range(5): + batched_policy_resp = _mock_batched_policy_ret(env_group.get_action_spaces(), batched_env_ret.ids()) + buffer.add_batched_transition(batched_env_ret, batched_policy_resp) + batched_env_ret, _ = env_group.step(batched_policy_resp) + + batched_policy_resp = _mock_batched_policy_ret(env_group.get_action_spaces(), batched_env_ret.ids()) + buffer.add_batched_transition(batched_env_ret, batched_policy_resp) + + assert len(buffer) == 0 + + buffer.truncate_episodes(batched_policy_resp.ids()) + assert len(buffer) == 5 * len(env_group) + assert len(buffer) == 5 * len(env_group) diff --git a/tests/integration/buffer/test_rollout_buffer_with_env.py b/tests/integration/buffer/test_rollout_buffer_with_env.py new file mode 100644 index 0000000..3f8d3ce --- /dev/null +++ b/tests/integration/buffer/test_rollout_buffer_with_env.py @@ -0,0 +1,106 @@ +import os + +import pytest + +from rlightning.buffer.utils.preprocessors import get_preprocessor_cls +from rlightning.env import EnvGroup +from rlightning.types import PolicyResponse +from rlightning.utils.builders import build_data_buffer +from rlightning.utils.config import BufferConfig +from rlightning.utils.ray import TaskSubmitter, resolve_object +from rlightning.utils.utils import InternalFlag +from tests.test_utils import _make_env_config, require_env_backend +from tensordict import TensorDict + +pytestmark = pytest.mark.integration + + +@pytest.fixture(scope="module", autouse=True) +def _auto_ray_cluster(): + # Ensure a Ray cluster is available for this module's tests + from tests.test_utils import setup_ray_cluster, teardown_ray_cluster + + setup_ray_cluster() + yield + teardown_ray_cluster() + + +def _make_buffer_config(capacity: int = 100) -> BufferConfig: + return BufferConfig.from_dict( + { + "type": "RolloutBuffer", + "capacity": capacity, + "storage": { + "type": "unified", + "device": "cpu", + }, + "sampler": { + "type": "all", + }, + } + ) + + +def _make_env(env_cfg, as_remote): + require_env_backend(env_cfg.backend) + os.environ["RLIGHTNING_REMOTE_ENV"] = "1" if as_remote else "0" + assert InternalFlag.REMOTE_ENV == as_remote + env, _ = EnvGroup.build_env(env_cfg) + if as_remote: + env_meta = env.init.remote() + else: + env_meta = env.init() + return env, env_meta + + +def _make_buffer(buffer_cfg, env_meta_list=None, as_remote: bool = False): + os.environ["RLIGHTNING_REMOTE_STORAGE"] = "1" if as_remote else "0" + assert InternalFlag.REMOTE_STORAGE == as_remote + + env_meta = env_meta_list[0] if env_meta_list is not None else None + if env_meta is not None and "observation_space" in env_meta: + observation_space = env_meta.observation_space + obs_preprocessor = get_preprocessor_cls(observation_space)(observation_space) + buffer = build_data_buffer(buffer_cls=buffer_cfg.type, buffer_cfg=buffer_cfg, obs_preprocessor=obs_preprocessor) + else: + buffer = build_data_buffer(buffer_cls=buffer_cfg.type, buffer_cfg=buffer_cfg) + + buffer.init(env_meta_list) + + return buffer + + +@pytest.mark.parametrize("env_as_remote", [False, True]) +@pytest.mark.parametrize("buffer_as_remote", [False, True]) +def test_sample(env_as_remote, buffer_as_remote): + require_env_backend("mujoco") + env_backend = "mujoco" + task_name = "Ant-v5" + num_env_workers = 1 + env_cfg = _make_env_config(env_backend, task_name, num_env_workers) + + env, _ = _make_env(env_cfg, env_as_remote) + buffer_cfg = _make_buffer_config(capacity=10) + buffer = _make_buffer(buffer_cfg, as_remote=buffer_as_remote) + + task_submitter = TaskSubmitter() + + env_ret = task_submitter.submit(env.reset, _block=True) + for _ in range(5): + action = task_submitter.submit(env.get_action_space, _block=True).sample() + policy_resp = PolicyResponse(env_id=env_ret.env_id, action=action, log_prob=0.0, entropy=0.0) + buffer.add_transition(env_ret.env_id, env_ret, policy_resp) + env_ret = task_submitter.submit(env.step, policy_resp, _block=True) + + action = task_submitter.submit(env.get_action_space, _block=True).sample() + policy_resp = PolicyResponse(env_id=env_ret.env_id, action=action, log_prob=0.0, entropy=0.0) + buffer.add_transition(env_ret.env_id, env_ret, policy_resp) + + buffer.truncate_one_episode(policy_resp.env_id) + assert len(buffer) == 5 + sample_data = buffer.sample(batch_size=5) + data = resolve_object(sample_data[0]) + data = TensorDict.from_dict(data, auto_batch_size=True) + assert len(data) == 5 + # after RolloutBuffer.sample, the buffer will clear automatically + assert len(buffer) == 0 diff --git a/tests/integration/buffer/test_sharded_buffer.py b/tests/integration/buffer/test_sharded_buffer.py new file mode 100644 index 0000000..4d09a0d --- /dev/null +++ b/tests/integration/buffer/test_sharded_buffer.py @@ -0,0 +1,206 @@ +import os + +import pytest +import ray +from tensordict import TensorDict + +from rlightning.types import BatchedData, PolicyResponse +from rlightning.utils.builders import build_data_buffer, build_env_group +from rlightning.utils.config import BufferConfig, ClusterConfig +from rlightning.utils.placement import GlobalResourceManager +from rlightning.utils.placement.placement_strategies import ( + PLACEMENT_STRATEGIES, + DefaultPlacementStrategy, +) +from rlightning.utils.placement.scheduling import ComponentScheduling, Scheduling +from rlightning.utils.ray import resolve_object +from tests.test_utils import _make_env_config, require_env_backend, setup_ray_cluster, teardown_ray_cluster + +pytestmark = pytest.mark.integration + + +@pytest.fixture(scope="module", autouse=True) +def _auto_ray_cluster(): + setup_ray_cluster() + yield + teardown_ray_cluster() + + +class _ShardedPlacementStrategy(DefaultPlacementStrategy): + """ + Placement strategy for testing sharded buffer on single node. + It allows placing multiple shards on the same node (round-robin). + """ + + def create_placement_groups(self): + num_buffer_storages = self.scheduling.buffer_worker.worker_num + if num_buffer_storages > 1: + node_ids = list(self._node_info["node_id_to_resources"].keys()) + + for storage_index in range(num_buffer_storages): + node_id = node_ids[storage_index % len(node_ids)] + self.buffer_strategies.append( + ray.util.scheduling_strategies.NodeAffinitySchedulingStrategy( + node_id=node_id, + soft=False, + ) + ) + return {} + + +# TODO: now only test sharded storage of ReplayBuffer type +def _make_sharded_buffer_config(capacity: int = 100, num_shards: int = 2) -> BufferConfig: + os.environ["RLIGHTNING_REMOTE_STORAGE"] = "1" + return BufferConfig.from_dict( + { + "type": "ReplayBuffer", + "capacity": capacity, + "storage": { + "type": "sharded", + "device": "cpu", + }, + "sampler": { + "type": "uniform", + }, + } + ) + + +def _init_global_resource_manager(num_shards=2, num_train_workers=2): + if "test_sharded" not in PLACEMENT_STRATEGIES: + PLACEMENT_STRATEGIES["test_sharded"] = _ShardedPlacementStrategy + + global_resource_manager = GlobalResourceManager.get_instance() + placement_cfg = ClusterConfig() + placement_cfg.placement.strategy = "test_sharded" + + scheduling = ComponentScheduling( + env_worker=[Scheduling(worker_num=1, num_cpus=0, num_gpus=0)], + train_worker=Scheduling(worker_num=num_train_workers, num_cpus=0, num_gpus=0), + eval_worker=Scheduling(worker_num=1, num_cpus=0, num_gpus=0), + buffer_worker=Scheduling(worker_num=num_shards, num_cpus=0, num_gpus=0), + ) + + global_resource_manager.initialize(placement_cfg, scheduling) + return global_resource_manager + + +def _reset_global_resource_manager(): + global_resource_manager = GlobalResourceManager.get_instance() + global_resource_manager.reset() + + +def _mock_batched_policy_ret(action_spaces, env_ids): + action_list = [resolve_object(action_space).sample() for action_space in action_spaces] + policy_resp_list = [ + PolicyResponse(env_id=env_id, action=action, log_prob=0.0, entropy=0.0) + for action, env_id in zip(action_list, env_ids) + ] + batched_policy_resp = BatchedData(env_ids, policy_resp_list) + return batched_policy_resp + + +def test_sharded_init(): + num_shards = 2 + num_train_workers = 2 + _init_global_resource_manager(num_shards=num_shards, num_train_workers=num_train_workers) + buffer_cfg = _make_sharded_buffer_config(num_shards=num_shards) + env_ids = [f"env_{i}" for i in range(4)] + + buffer = build_data_buffer(buffer_cls=buffer_cfg.type, buffer_cfg=buffer_cfg) + buffer.init(env_meta_list=None, env_ids=env_ids) + + assert len(buffer.storages) == num_shards + assert buffer.size() == 0 + + _reset_global_resource_manager() + + +@pytest.mark.parametrize("num_env_workers", [2, 4]) +@pytest.mark.parametrize("num_train_workers", [2, 4]) +@pytest.mark.parametrize("num_shards", [2]) +@pytest.mark.parametrize("steps", [10]) +def test_sharded_buffer(num_env_workers, num_train_workers, num_shards, steps): + require_env_backend("mujoco") + _init_global_resource_manager(num_shards=num_shards, num_train_workers=num_train_workers) + + buffer_cfg = _make_sharded_buffer_config(capacity=100, num_shards=num_shards) + + env_cfg = _make_env_config(env_backend="mujoco", task_name="Ant-v5", num_workers=num_env_workers) + env_group = build_env_group([env_cfg]) + batched_env_ret, _ = env_group.reset() + env_ids = batched_env_ret.ids() + + buffer = build_data_buffer(buffer_cls=buffer_cfg.type, buffer_cfg=buffer_cfg) + buffer.init(env_meta_list=None, env_ids=env_ids) + + # Add transitions + for _ in range(steps): + batched_policy_resp = _mock_batched_policy_ret(env_group.get_action_spaces(), batched_env_ret.ids()) + buffer.add_batched_transition(batched_env_ret, batched_policy_resp) + batched_env_ret, _ = env_group.step(batched_policy_resp) + + # Truncate episodes to flush to storage + truncations = [True] * len(batched_env_ret) + batched_policy_resp = _mock_batched_policy_ret(env_group.get_action_spaces(), batched_env_ret.ids()) + buffer.add_batched_transition(batched_env_ret, batched_policy_resp, truncations=truncations) + + # Check total size + expected_size = num_env_workers * steps + assert buffer.size() == expected_size + + # Test sampling + batch_size = num_env_workers * steps + sample_data = buffer.sample(batch_size=batch_size) + + assert len(sample_data) == num_train_workers + + total_samples = 0 + for data in sample_data: + data = resolve_object(data) + data = TensorDict.from_dict(data, auto_batch_size=True) + total_samples += len(data) + assert total_samples == batch_size + + _reset_global_resource_manager() + + +@pytest.mark.parametrize("num_env_workers", [2, 4]) +@pytest.mark.parametrize("num_train_workers", [2, 4]) +@pytest.mark.parametrize("num_shards", [2]) +@pytest.mark.parametrize("steps", [10]) +def test_add_data_async(num_env_workers, num_train_workers, num_shards, steps): + require_env_backend("mujoco") + _init_global_resource_manager(num_shards=num_shards, num_train_workers=num_train_workers) + + buffer_cfg = _make_sharded_buffer_config(capacity=100, num_shards=num_shards) + + env_cfg = _make_env_config(env_backend="mujoco", task_name="Ant-v5", num_workers=num_env_workers) + env_group = build_env_group([env_cfg]) + batched_env_ret, _ = env_group.reset() + env_ids = batched_env_ret.ids() + + buffer = build_data_buffer(buffer_cls=buffer_cfg.type, buffer_cfg=buffer_cfg) + buffer.init(env_meta_list=None, env_ids=env_ids) + + # Add transitions + for _ in range(steps): + batched_policy_resp = _mock_batched_policy_ret(env_group.get_action_spaces(), batched_env_ret.ids()) + buffer.add_batched_data_async(batched_env_ret) + buffer.add_batched_data_async(batched_policy_resp) + + batched_env_ret, _ = env_group.step(batched_policy_resp) + + assert buffer.size() == 0 + + # Explicitly truncate + batched_policy_resp = _mock_batched_policy_ret(env_group.get_action_spaces(), batched_env_ret.ids()) + buffer.add_batched_data_async(batched_env_ret) + buffer.add_batched_data_async(batched_policy_resp) + buffer.truncate_episodes(env_ids) + + # Check total size + expected_size = num_env_workers * steps + assert buffer.size() == expected_size + + _reset_global_resource_manager() diff --git a/tests/integration/env/__init__.py b/tests/integration/env/__init__.py new file mode 100644 index 0000000..b580d19 --- /dev/null +++ b/tests/integration/env/__init__.py @@ -0,0 +1 @@ +"""Environment integration tests.""" diff --git a/tests/integration/env/test_basic_ale.py b/tests/integration/env/test_basic_ale.py new file mode 100644 index 0000000..55857f6 --- /dev/null +++ b/tests/integration/env/test_basic_ale.py @@ -0,0 +1,144 @@ +"""This file implements test cases for basic ale tasks, which are integrated in gymnasium""" + +import os + +import pytest +import ray + +pytest.importorskip("ale_py") + +from rlightning.env.ale_env import ALEEnv +from rlightning.types import BatchedData, PolicyResponse +from rlightning.utils.builders import build_env_group +from rlightning.utils.config import EnvConfig +from rlightning.utils.utils import InternalFlag + +pytestmark = pytest.mark.integration + +@pytest.fixture(scope="module", autouse=True) +def _auto_ray_cluster(): + # Ensure a Ray cluster is available for this module's tests + from tests.test_utils import setup_ray_cluster, teardown_ray_cluster + + setup_ray_cluster(num_gpus=1) + yield + teardown_ray_cluster() + + +@pytest.mark.parametrize("task", ["ALE/Adventure-v5", "ALE/Alien-v5"]) +def test_env_logic(task: str): + + max_episode_steps = 10 + + config = EnvConfig(name=task, task=task, backend="ale", max_episode_steps=max_episode_steps) + + env = ALEEnv(config) + + _ = env.reset() + + terminated = truncated = False + step_cnt = 0 + + while not (terminated or truncated): + + action = env.action_space.sample() + policy_resp = PolicyResponse(env_id=env.env_id, action=action) + env_ret = env.step(policy_resp) + observation, reward, terminated, truncated, info = ( + env_ret.observation, + env_ret.last_reward, + env_ret.last_terminated, + env_ret.last_truncated, + env_ret.info, + ) + + step_cnt += 1 + if not terminated and step_cnt >= max_episode_steps: + assert truncated, (step_cnt, max_episode_steps, terminated, truncated) + + assert step_cnt <= max_episode_steps, ( + step_cnt, + max_episode_steps, + terminated, + truncated, + ) + + +@pytest.mark.parametrize("task", ["ALE/Adventure-v5", "ALE/Alien-v5"]) +@pytest.mark.parametrize("env_as_remote", [True, False]) +def test_feed_into_env_group(env_as_remote: bool, task: str): + + max_episode_steps = 100 + + env_cfgs = [ + EnvConfig( + name=task, + backend="ale", + task=task, + num_workers=1, + num_envs=1, + max_episode_steps=max_episode_steps, + env_parameters=dict(), + ) + ] + os.environ["RLIGHTNING_REMOTE_ENV"] = "1" if env_as_remote else "0" + assert InternalFlag.REMOTE_ENV == env_as_remote + + env_group = build_env_group(env_cfgs) + + assert len(env_group) == 1 + + batched_env_ret, truncations = env_group.reset() + + assert len(batched_env_ret) == 1 + + if env_as_remote: + env_id, env_ret_future = list(batched_env_ret.items())[0] + assert isinstance(env_ret_future, ray.ObjectRef), type(env_ret_future) + env_ret = ray.get(env_ret_future) + else: + env_id, env_ret = list(batched_env_ret.items())[0] + + terminated = truncated = False + step_cnt = 0 + + while not (terminated or truncated): + + action_spaces = env_group.get_action_spaces() + if env_as_remote: + actions = [ray.get(e).sample() for e in action_spaces] + else: + actions = [e.sample() for e in action_spaces] + + policy_resp_list = [ + PolicyResponse(env_id=env_group.env_ids[i], action=action) + for i, action in enumerate(actions) + ] + + assert len(policy_resp_list) == 1, len(actions) + policy_rets = BatchedData( + ids=[env_id], + data=policy_resp_list, + ) + env_rets, truncations = env_group.step(policy_rets) + + if env_as_remote: + assert len(env_rets) == 1 + env_ret_future = env_rets.values()[0] + env_ret = ray.get(env_ret_future) + else: + env_ret = env_rets.values()[0] + # observation, reward, terminated, truncated, info = env_rets[0] + + terminated, truncated = env_ret.last_terminated, env_ret.last_truncated + + step_cnt += 1 + if not terminated and step_cnt >= max_episode_steps: + assert truncated, (step_cnt, max_episode_steps, terminated, truncated) + + assert step_cnt <= max_episode_steps, ( + step_cnt, + max_episode_steps, + terminated, + truncated, + ) diff --git a/tests/integration/env/test_basic_mujoco.py b/tests/integration/env/test_basic_mujoco.py new file mode 100644 index 0000000..088de36 --- /dev/null +++ b/tests/integration/env/test_basic_mujoco.py @@ -0,0 +1,145 @@ +"""This file implements test cases for basic mujoco tasks, which are integrated in gymnasium""" + +import os + +import pytest +import ray + +pytest.importorskip("mujoco") + +from rlightning.env.mujoco_env import MujocoEnv +from rlightning.types import BatchedData, PolicyResponse +from rlightning.utils.builders import build_env_group +from rlightning.utils.config import EnvConfig +from rlightning.utils.utils import InternalFlag + +pytestmark = pytest.mark.integration + + +@pytest.fixture(scope="module", autouse=True) +def _auto_ray_cluster(): + # Ensure a Ray cluster is available for this module's tests + from tests.test_utils import setup_ray_cluster, teardown_ray_cluster + + setup_ray_cluster(num_gpus=1) + yield + teardown_ray_cluster() + + +@pytest.mark.parametrize("task", ["Ant-v5", "HalfCheetah-v5", "Hopper-v5", "Humanoid-v5"]) +def test_env_logic(task: str): + + max_episode_steps = 10 + + config = EnvConfig(name=task, task=task, backend="mujoco", max_episode_steps=max_episode_steps) + + env = MujocoEnv(config) + + _ = env.reset() + + terminated = truncated = False + step_cnt = 0 + + while not (terminated or truncated): + + action = env.action_space.sample() + policy_resp = PolicyResponse(env_id=env.env_id, action=action) + env_ret = env.step(policy_resp) + observation, reward, terminated, truncated, info = ( + env_ret.observation, + env_ret.last_reward, + env_ret.last_terminated, + env_ret.last_truncated, + env_ret.info, + ) + + step_cnt += 1 + if not terminated and step_cnt >= max_episode_steps: + assert truncated, (step_cnt, max_episode_steps, terminated, truncated) + + assert step_cnt <= max_episode_steps, ( + step_cnt, + max_episode_steps, + terminated, + truncated, + ) + + +@pytest.mark.parametrize("task", ["Ant-v5", "HalfCheetah-v5", "Hopper-v5", "Humanoid-v5"]) +@pytest.mark.parametrize("env_as_remote", [True, False]) +def test_feed_into_env_client(env_as_remote: bool, task: str): + + max_episode_steps = 100 + + env_cfgs = [ + EnvConfig( + name=task, + backend="mujoco", + task=task, + num_workers=1, + num_envs=1, + max_episode_steps=max_episode_steps, + env_parameters=dict(), + ) + ] + os.environ["RLIGHTNING_REMOTE_ENV"] = "1" if env_as_remote else "0" + assert InternalFlag.REMOTE_ENV == env_as_remote + + env_group = build_env_group(env_cfgs) + + assert len(env_group) == 1 + + batched_env_ret, truncations = env_group.reset() + + assert len(batched_env_ret) == 1 + + if env_as_remote: + env_id, env_ret_future = list(batched_env_ret.items())[0] + assert isinstance(env_ret_future, ray.ObjectRef), type(env_ret_future) + env_ret = ray.get(env_ret_future) + else: + env_id, env_ret = list(batched_env_ret.items())[0] + + terminated = truncated = False + step_cnt = 0 + + while not (terminated or truncated): + + action_spaces = env_group.get_action_spaces() + if env_as_remote: + actions = [ray.get(e).sample() for e in action_spaces] + else: + actions = [e.sample() for e in action_spaces] + + policy_resp_list = [ + PolicyResponse(env_id=env_group.env_ids[i], action=action) + for i, action in enumerate(actions) + ] + + assert len(policy_resp_list) == 1, len(actions) + policy_rets = BatchedData( + ids=[env_id], + data=policy_resp_list, + ) + env_rets, truncations = env_group.step(policy_rets) + + if env_as_remote: + assert len(env_rets) == 1 + env_ret_future = env_rets.values()[0] + env_ret = ray.get(env_ret_future) + else: + env_ret = env_rets.values()[0] + # observation, reward, terminated, truncated, info = env_rets[0] + + terminated, truncated = env_ret.last_terminated, env_ret.last_truncated + + step_cnt += 1 + if not terminated and step_cnt >= max_episode_steps: + assert truncated, (step_cnt, max_episode_steps, terminated, truncated) + + assert step_cnt <= max_episode_steps, ( + step_cnt, + max_episode_steps, + terminated, + truncated, + ) diff --git a/tests/integration/env/test_env_group.py b/tests/integration/env/test_env_group.py new file mode 100644 index 0000000..61d804b --- /dev/null +++ b/tests/integration/env/test_env_group.py @@ -0,0 +1,316 @@ +import os +from pathlib import Path + +import hydra +import pytest +import torch + +from hydra import compose, initialize_config_dir +from omegaconf import DictConfig + +from rlightning.types import BatchedData, PolicyResponse +from rlightning.utils.builders import build_env_group +from rlightning.utils.ray import resolve_object + +# set parent dir for examples +# Resolve repo root even after moving tests into tests/integration/... +EXAMPLES_DIR = Path(__file__).resolve().parents[3] / "examples" +from rlightning.utils.utils import InternalFlag +from tests.test_utils import _make_env_config, require_env_backend + +pytestmark = pytest.mark.integration + + +@pytest.fixture(scope="module", autouse=True) +def _auto_ray_cluster(): + # Ensure a Ray cluster is available for this module's tests + from tests.test_utils import setup_ray_cluster, teardown_ray_cluster + + setup_ray_cluster(num_cpus=4, num_gpus=2) + yield + teardown_ray_cluster() + + +@pytest.fixture +def isaac_cleanup_registry(scope="module"): + """ + Fixture that provides a list of instances containing isaaclab for cleanup. + The creation happens in the test, but cleanup happens here. + Used to avoid clear failure caused by isaaclab in testing. + """ + registry = [] + yield registry + + # Teardown: runs after the test finishes + print(f"Fixture teardown: Closing {len(registry)} registered...") + for instance in registry: + try: + if hasattr(instance, "close"): + instance.close() + else: + print("Warning: instance does not have a close() method.") + except Exception as e: + print(f"Warning: Error closing instance: {e}") + print("Fixture teardown: All registered instance closed.") + + +def _mock_batched_policy_ret(action_spaces, env_ids): + action_list = [resolve_object(action_space).sample() for action_space in action_spaces] + policy_resp_list = [ + PolicyResponse(env_id=env_id, action=action, log_prob=0.0, entropy=0.0) + for action, env_id in zip(action_list, env_ids) + ] + batched_policy_resp = BatchedData(env_ids, policy_resp_list) + return batched_policy_resp + + +def _mock_policy_ret(action_space, env_id): + action = action_space.sample() + policy_resp = PolicyResponse(env_id=env_id, action=action, log_prob=0.0, entropy=0.0) + return policy_resp + + +# # to be fixed: it can pass individually but fails in `make test-unit` +# @pytest.mark.parametrize( +# "config_path,config_name", +# [ +# ("rslrl_isaaclab/conf/env", "isaaclab_mimic"), +# ], +# ) +# def test_hydra_entrypoint(config_path: str, config_name: str, isaac_cleanup_registry): +# """Test hydra entrypoint for building env group from config.""" +# # set parent dir for examples +# EXAMPLES_DIR = Path(__file__).parent.parent.parent / "examples" + +# @hydra.main() +# def entrypoint(cfg: DictConfig): +# print(f"--- Complete Config ---\n{cfg}") +# env_cfg = EnvConfig.from_omegaconf(cfg) +# env_group = build_env_group([env_cfg]) +# isaac_cleanup_registry.append(env_group) + +# with initialize_config_dir(config_dir=str(EXAMPLES_DIR / config_path)): +# cfg = compose(config_name=config_name) +# entrypoint(cfg) + + +@pytest.mark.parametrize("as_remote", [True, False]) +@pytest.mark.parametrize("task", ["Ant-v5", "HalfCheetah-v5", "Hopper-v5", "Humanoid-v5"]) +def test_mujoco_make_env_group(task, as_remote): + require_env_backend("mujoco") + env_cfgs = [ + _make_env_config(env_backend="mujoco", task_name=task, num_workers=2), + ] + os.environ["RLIGHTNING_REMOTE_ENV"] = "1" if as_remote else "0" + assert InternalFlag.REMOTE_ENV == as_remote + env_group = build_env_group(env_cfgs) + + assert len(env_group) == 2 + + +# FIXME: it can pass individually but fails in `make test-unit` +# @pytest.mark.parametrize("task", ["Tracking-Flat-G1-v0"]) +# @pytest.mark.parametrize("num_envs", [1]) +# def test_isaaclab_manager_based_env_group(task: str, num_envs: int, isaac_cleanup_registry): +# task = "Tracking-Flat-G1-v0" + +# max_episode_steps = 10 + +# config = EnvConfig( +# name=task, +# task=task, +# num_envs=num_envs, +# backend="isaac_manager_based", +# max_episode_steps=max_episode_steps, +# env_kwargs=dict( +# device=f"cuda:{torch.cuda.current_device()}", +# headless=True, +# commands=dict(motion={"motion_file": "tests/data/fallAndGetUp1_subject1.npz"}), +# ), +# ) + +# env_cfgs = [config] + +# # env_cfgs = [ +# # _make_env_config( +# # env_backend="isaac_manager_based", +# # task=task, +# # num_envs=num_envs, +# # ) +# # ] + +# env_group = build_env_group(env_cfgs) +# assert len(env_group) == 1 + +# isaac_cleanup_registry.append(env_group) + + +@pytest.mark.parametrize("as_remote", [True, False]) +def test_env_group_shape_check_fail_case(as_remote): + """ + Test that the environment group raises a ValueError when environments have mismatched + observation or action space shapes. + """ + require_env_backend("mujoco") + env_cfgs = [ + _make_env_config(env_backend="mujoco", task_name="Ant-v5", num_workers=1), + _make_env_config(env_backend="mujoco", task_name="Hopper-v5", num_workers=1), + ] + os.environ["RLIGHTNING_REMOTE_ENV"] = "1" if as_remote else "0" + assert InternalFlag.REMOTE_ENV == as_remote + + env_group = build_env_group(env_cfgs) + with pytest.raises(ValueError): + env_group.init() + + +@pytest.mark.parametrize("as_remote", [True, False]) +def test_env_group_shape_check_success_case(as_remote): + require_env_backend("mujoco") + env_cfgs = [ + _make_env_config(env_backend="mujoco", task_name="Ant-v5", num_workers=1), + _make_env_config(env_backend="mujoco", task_name="Ant-v5", num_workers=1), + ] + + os.environ["RLIGHTNING_REMOTE_ENV"] = "1" if as_remote else "0" + assert InternalFlag.REMOTE_ENV == as_remote + + env_group = build_env_group(env_cfgs) + env_group.init() + + +@pytest.mark.parametrize("as_remote", [True, False]) +@pytest.mark.parametrize("num_env_workers", [1, 2, 4]) +def test_env_sync_step(as_remote, num_env_workers): + require_env_backend("mujoco") + env_cfgs = [ + _make_env_config(env_backend="mujoco", task_name="Ant-v5", num_workers=num_env_workers), + ] + + os.environ["RLIGHTNING_REMOTE_ENV"] = "1" if as_remote else "0" + assert InternalFlag.REMOTE_ENV == as_remote + + env_group = build_env_group(env_cfgs) + env_meta_list = env_group.init() + + batched_env_ret, _ = env_group.reset() + assert isinstance(batched_env_ret, BatchedData) + assert len(batched_env_ret.ids()) == num_env_workers + for i, (env_id, env_ret) in enumerate(batched_env_ret.items()): + assert ( + resolve_object(env_ret).observation.shape + == resolve_object(env_meta_list[i]).observation_space.shape + ) + + for _ in range(5): + policy_resp_list = [ + _mock_policy_ret(resolve_object(env_meta).action_space, env_id) + for env_meta, env_id in zip(env_meta_list, batched_env_ret.ids()) + ] + batched_policy_resp = BatchedData(batched_env_ret.ids(), policy_resp_list) + batched_env_ret, _ = env_group.step(batched_policy_resp) + assert isinstance(batched_env_ret, BatchedData) + assert len(batched_env_ret.ids()) == num_env_workers + + +@pytest.mark.parametrize("as_remote", [True, False]) +@pytest.mark.parametrize("num_env_workers", [1, 2, 4]) +def test_env_step_counter(as_remote, num_env_workers): + require_env_backend("mujoco") + env_cfgs = [ + _make_env_config(env_backend="mujoco", task_name="Ant-v5", num_workers=num_env_workers), + ] + + os.environ["RLIGHTNING_REMOTE_ENV"] = "1" if as_remote else "0" + assert InternalFlag.REMOTE_ENV == as_remote + + env_group = build_env_group(env_cfgs) + env_meta_list = env_group.init() + + max_episode_steps = 2 + with env_group.auto_reset(max_episode_steps=max_episode_steps): + # before reset + for env_id in env_group.env_ids: + assert env_group.step_counter.get_steps(env_id) == -1 + assert not env_group.step_counter.is_reached_max_steps(env_id) + + batched_env_ret, truncations = env_group.reset() + # after reset + for env_id in env_group.env_ids: + assert env_group.step_counter.get_steps(env_id) == 0 + assert not env_group.step_counter.is_reached_max_steps(env_id) + for truncated in truncations: + assert not truncated + + policy_resp_list = [ + _mock_policy_ret(resolve_object(env_meta).action_space, env_id) + for env_meta, env_id in zip(env_meta_list, batched_env_ret.ids()) + ] + batched_policy_resp = BatchedData(batched_env_ret.ids(), policy_resp_list) + + batched_env_ret, truncations = env_group.step(batched_policy_resp) + # after 1st step + for env_id in env_group.env_ids: + assert env_group.step_counter.get_steps(env_id) == 1 + assert not env_group.step_counter.is_reached_max_steps(env_id) + for truncated in truncations: + assert not truncated + + batched_env_ret, truncations = env_group.step(batched_policy_resp) + # after 2nd step + for env_id in env_group.env_ids: + assert env_group.step_counter.get_steps(env_id) == max_episode_steps + assert env_group.step_counter.is_reached_max_steps(env_id) + for truncated in truncations: + assert truncated + + batched_env_ret, truncations = env_group.step(batched_policy_resp) + # after 3nd step (should auto reset) + for env_id in env_group.env_ids: + assert env_group.step_counter.get_steps(env_id) == 0 + assert not env_group.step_counter.is_reached_max_steps(env_id) + for truncated in truncations: + assert not truncated + + +@pytest.mark.parametrize("as_remote", [True, False]) +@pytest.mark.parametrize("num_env_workers", [1, 2, 4]) +def test_auto_reset_sync(as_remote, num_env_workers): + require_env_backend("mujoco") + env_cfgs = [ + _make_env_config(env_backend="mujoco", task_name="Ant-v5", num_workers=num_env_workers), + ] + + os.environ["RLIGHTNING_REMOTE_ENV"] = "1" if as_remote else "0" + assert InternalFlag.REMOTE_ENV == as_remote + + env_group = build_env_group(env_cfgs) + env_meta_list = env_group.init() + + max_episode_steps = 5 + with env_group.auto_reset(max_episode_steps=max_episode_steps): + # test reset + batched_env_ret, truncations = env_group.reset() + for truncated in truncations: + assert not truncated + for env_id in env_group.env_ids: + assert env_group.step_counter[env_id] == 0 + + # test step + for step in range(1, 20): + policy_resp_list = [ + _mock_policy_ret(resolve_object(env_meta).action_space, env_id) + for env_meta, env_id in zip(env_meta_list, batched_env_ret.ids()) + ] + batched_policy_resp = BatchedData(batched_env_ret.ids(), policy_resp_list) + batched_env_ret, truncations = env_group.step(batched_policy_resp) + + expected_step = step % (max_episode_steps + 1) + + for env_id in env_group.env_ids: + assert env_group.step_counter[env_id] == expected_step + for truncated in truncations: + if expected_step == max_episode_steps: + assert truncated + else: + assert not truncated diff --git a/tests/integration/env/test_env_group_performance.py b/tests/integration/env/test_env_group_performance.py new file mode 100644 index 0000000..b7ed8cf --- /dev/null +++ b/tests/integration/env/test_env_group_performance.py @@ -0,0 +1,129 @@ +import pytest + +try: + import gymnasium as gym +except Exception: + gym = None + +import numpy as np + +from rlightning.env.base_env import BaseEnv +from rlightning.types import EnvRet +from rlightning.utils.builders import build_env_group, build_policy_group +from rlightning.utils.config import EnvConfig, PolicyConfig, ClusterConfig +from rlightning.utils.registry import ENVS, POLICIES + + +class _DummyInnerEnv: + def __init__(self): + self.observation_space = gym.spaces.Box(low=0.0, high=1.0, shape=(4,), dtype=np.float32) + self.action_space = gym.spaces.Box(low=-1.0, high=1.0, shape=(2,), dtype=np.float32) + + def reset(self, *args, **kwargs): + return np.zeros(self.observation_space.shape, dtype=self.observation_space.dtype), {} + + def step(self, action): + observation = np.zeros(self.observation_space.shape, dtype=self.observation_space.dtype) + reward = 0.0 + terminated = False + truncated = False + info = {} + return observation, reward, terminated, truncated, info + + +class DummyPerfEnv(BaseEnv): + def __init__(self, config, preprocess_fn, **kwargs): + super().__init__(config, preprocess_fn) + self.env = _DummyInnerEnv() + self.observation_space = self.env.observation_space + self.action_space = self.env.action_space + + def reset(self, *args, **kwargs): + observation, _info = self.env.reset(*args, **kwargs) + return EnvRet(env_id=self.env_id, observation=observation) + + def step(self, policy_resp): + action = self._preprocess_fn(policy_resp) if self._preprocess_fn else policy_resp + observation, reward, terminated, truncated, info = self.env.step(action) + return EnvRet( + env_id=self.env_id, + observation=observation, + last_reward=reward, + last_terminated=terminated, + last_truncated=truncated, + info=info, + ) + + +class DummyPolicy: + def __init__(self, config, role_type=None): + self.config = config + self.role_type = role_type + + def init_weight_buffer(self): + return None + + +@pytest.mark.integration +@pytest.mark.slow +def test_env_group_worker_perf(performance_timer, monkeypatch): + if gym is None: + pytest.skip("gymnasium not available") + + env_key = "ale" + policy_key = "mock_perf_policy" + + monkeypatch.setitem(ENVS.module_dict, env_key, DummyPerfEnv) + monkeypatch.setitem(POLICIES.module_dict, policy_key, DummyPolicy) + + env_cfg = EnvConfig.load_yaml("examples/wbc_tracking/conf/env/single_task.yaml") + env_cfg.backend = env_key + env_cfg.task = "DummyTask-v0" + env_cfg.num_envs = 1 + + policy_cfg = PolicyConfig(type=policy_key) + + worker_counts = [1, 2, 4] + results = [] + + for num_workers in worker_counts: + env_cfg.num_workers = num_workers + + performance_timer.start() + env_group = build_env_group(env_cfg) + env_time_s = performance_timer.stop() + + performance_timer.start() + policy_group = build_policy_group( + policy_key, + policy_cfg, + cluster_cfg=ClusterConfig(), + ) + policy_time_s = performance_timer.stop() + + env_group.close() + _ = policy_group + + env_time_s = max(env_time_s, 1e-9) + policy_time_s = max(policy_time_s, 1e-9) + + results.append( + { + "num_workers": num_workers, + "env_time_s": env_time_s, + "env_throughput": num_workers / env_time_s, + "policy_time_s": policy_time_s, + "policy_throughput": 2 / policy_time_s, + } + ) + print( + f"workers={num_workers} " + f"env_time_s={env_time_s:.6f} env_throughput={num_workers / env_time_s:.2f} " + f"policy_time_s={policy_time_s:.6f} policy_throughput={2 / policy_time_s:.2f}" + ) + + for metrics in results: + assert metrics["env_time_s"] >= 0.0 + assert metrics["policy_time_s"] >= 0.0 + assert metrics["env_throughput"] > 0.0 + assert metrics["policy_throughput"] > 0.0 diff --git a/tests/integration/env/test_isaac_manager_based.py b/tests/integration/env/test_isaac_manager_based.py new file mode 100644 index 0000000..45025ec --- /dev/null +++ b/tests/integration/env/test_isaac_manager_based.py @@ -0,0 +1,123 @@ +# FIXME: it can pass individually but fails in `make test-isaaclab` + +"""This file implemens test cases for tasks registered in isaac_marl""" + +import sys +from pathlib import Path + +import pytest +import torch + +pytest.importorskip("isaaclab") + +from rlightning.env.isaac_env import IsaacManagerBasedRLEnv +from rlightning.types import PolicyResponse +from rlightning.utils.config import EnvConfig + +pytestmark = [pytest.mark.integration, pytest.mark.isaaclab, pytest.mark.gpu] + + +@pytest.fixture(scope="module", autouse=True) +def _auto_ray_cluster(): + # Ensure a Ray cluster is available for this module's tests + from tests.test_utils import setup_ray_cluster, teardown_ray_cluster + + setup_ray_cluster(num_gpus=1) + yield + teardown_ray_cluster() + + +@pytest.fixture +def env_instance(): + """Fixture to manage the lifecycle of the environment.""" + env = _make_env() + yield env + + # Teardown: This code runs after the test function finishes + print("Fixture teardown: Closing env...") + try: + env.close() + except Exception as e: + print(f"Warning: Error closing env: {e}") + print("Fixture teardown: Env closed.") + + +def _make_env(): + task = "Tracking-Flat-G1-v0" + + max_episode_steps = 10 + root = Path(__file__).resolve().parents[3] + motion_dir = root / ".data" / "lafan1" / "retargeted" / "wbc_tracking" + examples_dir = root / "examples" + + if str(examples_dir) not in sys.path: + sys.path.insert(0, str(examples_dir)) + + # The tracking example pulls in the humanoid retargeting stack, which is + # optional and not installed by the plain `isaaclab` extra. + pytest.importorskip("mink") + if not motion_dir.exists(): + pytest.skip(f"WBC motion asset directory not found: {motion_dir}") + if not any(motion_dir.glob("*.npz")): + pytest.skip(f"No WBC motion files found under: {motion_dir}") + + config = EnvConfig( + name=task, + task=task, + backend="isaac_manager_based", + max_episode_steps=max_episode_steps, + env_kwargs=dict( + env_spec="wbc_tracking.envs", + launcher=dict(headless=True), + env_cfg=dict( + module="wbc_tracking.envs.flat_env_cfg::G1FlatEnvCfg", + override=dict( + commands=dict( + motion={ + "motion_dir": str(motion_dir), + } + ) + ), + ), + ), + ) + + original_argv = sys.argv[:] + # Clear the parameters to prevent Isaac Sim from crashing due to reading + # the -m parameter of pytest. + sys.argv = [sys.argv[0]] + + try: + env = IsaacManagerBasedRLEnv(config) + finally: + sys.argv = original_argv + return env + + +def test_env_functionalities(env_instance): + env = env_instance + env_ret = env.reset() + + for step_cnt in range(100): + + action = torch.tensor(env.get_action_space().sample()) + policy_resp = PolicyResponse(env_id=env.env_id, action=action) + env_ret = env.step(policy_resp) + + observation, reward, terminated, truncated, info = ( + env_ret.observation, + env_ret.last_reward, + env_ret.last_terminated, + env_ret.last_truncated, + env_ret.info, + ) + + print("reward:", reward.shape, reward.mean()) + print( + "step:", + step_cnt, + "terminated envs:", + terminated.sum(), + "truncated envs:", + truncated.sum(), + ) diff --git a/tests/integration/env/test_remote_env.py b/tests/integration/env/test_remote_env.py new file mode 100644 index 0000000..947f81f --- /dev/null +++ b/tests/integration/env/test_remote_env.py @@ -0,0 +1,204 @@ +import os +import time +from typing import List + +import numpy as np +import pytest +import ray + +from rlightning.buffer.utils.preprocessors import get_preprocessor_cls +from rlightning.types import BatchedData, PolicyResponse +from rlightning.utils.builders import build_data_buffer, build_env_group +from rlightning.utils.config import BufferConfig, EnvConfig +from rlightning.utils.ray import TaskSubmitter, resolve_object +from rlightning.utils.utils import InternalFlag +from tests.tests_utils.envs.utils_remote_env import ( + EnvClientWorker, + EnvServerWorker, + MockPiperEnv, +) + +pytestmark = pytest.mark.integration + + +@pytest.fixture(scope="module", autouse=True) +def _auto_ray_cluster(): + # Ensure a Ray cluster is available for this module's tests + from tests.test_utils import setup_ray_cluster, teardown_ray_cluster + + setup_ray_cluster(num_cpus=4, num_gpus=2, local_mode=False, log_to_driver=True) + yield + teardown_ray_cluster() + + +def _mock_batched_policy_response(batched_env_ret: BatchedData, policy_as_remote=False) -> List[PolicyResponse]: + policy_response_list = [] + for env_ret in batched_env_ret.values(): + env_ret = resolve_object(env_ret) + action = np.random.uniform(-1, 1, size=(7,)) + policy_response = PolicyResponse(env_id=env_ret.env_id, action=action) + if policy_as_remote: + policy_response = ray.put(policy_response) + policy_response_list.append(policy_response) + + return BatchedData(batched_env_ret.ids(), policy_response_list) + + +@pytest.mark.parametrize("num_envs", [4]) +@pytest.mark.parametrize("num_steps", [10]) +def test_env_server_client_interaction(num_envs, num_steps): + server = EnvServerWorker.remote() + address, port = ray.get(server.get_address_port.remote()) + + success = server.run.remote(num_envs, num_steps, timeout=100) + workers = [EnvClientWorker.remote(address, port, num_steps) for _ in range(num_envs)] + _ = [worker.run.remote() for worker in workers] + + success = ray.get(success) + + assert success, "Env server-client interaction test failed." + + server.close.remote() + + +@pytest.mark.parametrize("as_remote", [False, True]) +@pytest.mark.parametrize("num_steps", [10]) +@pytest.mark.parametrize("num_envs", [4]) +def test_env_group_with_env_server_to_interaction(as_remote, num_steps, num_envs): + task_submitter = TaskSubmitter() + + os.environ["RLIGHTNING_REMOTE_ENV"] = "1" if as_remote else "0" + assert InternalFlag.REMOTE_ENV == as_remote + + env_server_cfg = EnvConfig( + name="test_PiperEnv_Server", + backend="env_server", + task="real_world", + zmq_port="6366", + ) + env_cfgs = [env_server_cfg] + env_group = build_env_group(env_cfgs) + _ = env_group.init() + + address, port = task_submitter.submit(env_group.env_servers[0].get_address_port, _block=True) + + workers = [EnvClientWorker.remote(address, port, num_steps) for _ in range(num_envs)] + _ = [worker.run.remote() for worker in workers] + + batched_env_ret, _ = env_group.reset() + # two times reset called defined in RemoteEnvClient.run() + expect_cnt = (1 + num_steps) * num_envs * 2 + cnt = len(batched_env_ret) + + step = 0 + while True: + step += 1 + batched_policy_resp = _mock_batched_policy_response(batched_env_ret, as_remote) + env_group.step_async(batched_policy_resp) + batched_env_ret, _ = env_group.collect_async() + cnt += len(batched_env_ret) + + print(f"Step {step}: received {len(batched_env_ret)}, total {cnt}.", flush=True) + # FIXME cannot detect the wrong case that clients send more returns than expected + # and if not receive enough, the test will hang here forever. + if cnt == expect_cnt: + break + + env_group.close() + + +def _make_buffer_config(capacity: int = 100) -> BufferConfig: + return BufferConfig.from_dict( + { + "type": "ReplayBuffer", + "capacity": capacity, + "auto_truncate_episode": True, + "storage": { + "type": "unified", + "device": "cpu", + }, + "sampler": { + "type": "uniform", + }, + } + ) + + +def _make_buffer(buffer_cfg, env_meta_list=None, as_remote: bool = False): + os.environ["RLIGHTNING_REMOTE_STORAGE"] = "1" if as_remote else "0" + assert InternalFlag.REMOTE_STORAGE == as_remote + + env_meta = env_meta_list[0] if env_meta_list is not None else None + if env_meta is not None and "observation_space" in env_meta: + observation_space = env_meta.observation_space + obs_preprocessor = get_preprocessor_cls(observation_space)(observation_space) + buffer = build_data_buffer( + buffer_cls=buffer_cfg.type, + buffer_cfg=buffer_cfg, + obs_preprocessor=obs_preprocessor, + ) + else: + buffer = build_data_buffer( + buffer_cls=buffer_cfg.type, + buffer_cfg=buffer_cfg, + postprocess_fn=MockPiperEnv.episode_postprocess_fn, + ) + + buffer.init(env_meta_list) + + return buffer + + +@pytest.mark.parametrize("as_remote", [True, False]) +@pytest.mark.parametrize("num_steps", [10]) +@pytest.mark.parametrize("num_envs", [4]) +def test_rollout_and_collect_replay_buffer_with_remote_env(as_remote, num_steps, num_envs): + task_submitter = TaskSubmitter() + + os.environ["RLIGHTNING_REMOTE_ENV"] = "1" if as_remote else "0" + assert InternalFlag.REMOTE_ENV == as_remote + + buffer_cfg = _make_buffer_config(capacity=100) + buffer = _make_buffer(buffer_cfg, env_meta_list=None, as_remote=as_remote) + + env_server_cfg = EnvConfig( + name="test_PiperEnv_Server", + backend="env_server", + task="real_world", + zmq_port="6366", + ) + env_cfgs = [env_server_cfg] + env_group = build_env_group(env_cfgs) + _ = env_group.init() + address, port = task_submitter.submit(env_group.env_servers[0].get_address_port, _block=True) + + workers = [EnvClientWorker.remote(address, port, num_steps) for _ in range(num_envs)] + _ = [worker.run.remote() for worker in workers] + + # two times reset called defined in RemoteEnvClient.run() + expect_cnt = (1 + num_steps) * num_envs * 2 + + batched_env_ret, _ = env_group.reset() + cnt = len(batched_env_ret) + step = 0 + buffer.add_batched_data_async(batched_env_ret) + while True: + batched_policy_resp = _mock_batched_policy_response(batched_env_ret, as_remote) + buffer.add_batched_data_async(batched_policy_resp) + + env_group.step_async(batched_policy_resp) + batched_env_ret, _ = env_group.collect_async() + buffer.add_batched_data_async(batched_env_ret) + + step += 1 + cnt += len(batched_env_ret) + + print(f"Step {step}: received {len(batched_env_ret)}, total {cnt}.", flush=True) + # FIXME cannot detect the wrong case that clients send more returns than expected + # and if not receive enough, the test will hang here forever. + if cnt == expect_cnt: + break + + assert buffer.size() > 0 + + env_group.close() diff --git a/tests/integration/policy/__init__.py b/tests/integration/policy/__init__.py new file mode 100644 index 0000000..3469c24 --- /dev/null +++ b/tests/integration/policy/__init__.py @@ -0,0 +1 @@ +"""Policy integration tests.""" diff --git a/tests/integration/policy/test_ddp_checkpoint.py b/tests/integration/policy/test_ddp_checkpoint.py new file mode 100644 index 0000000..6588bca --- /dev/null +++ b/tests/integration/policy/test_ddp_checkpoint.py @@ -0,0 +1,140 @@ +import pytest +import torch + +from rlightning.utils.config import ( + ClusterConfig, + Config, + PolicyConfig, + TrainConfig, + WeightBufferConfig, +) +from rlightning.utils.utils import InternalFlag +from tests.test_utils import ( + get_test_ray_address, + init_test_ray, + setup_ray_cluster, + should_skip_gpu_test, + should_skip_ray_test, + teardown_ray_cluster, +) +from tests.tests_utils.ddp_checkpoint_policy import DDPCheckpointPolicy + + +@pytest.fixture(scope="module", autouse=True) +def _ddp_ray_cluster(): + if should_skip_ray_test(): + pytest.skip("Ray not available") + if should_skip_gpu_test(): + pytest.skip("CUDA not available") + if torch.cuda.device_count() < 2: + pytest.skip("DDP checkpoint test requires at least 2 GPUs") + + started_cluster = False + import ray + + ray_address = get_test_ray_address(default="local") + + if ray.is_initialized(): + ray.shutdown() + + if ray_address == "local": + setup_ray_cluster(num_cpus=2, num_gpus=2, local_mode=False) + started_cluster = True + else: + init_test_ray(log_to_driver=False, default_address=ray_address) + + if ray.cluster_resources().get("GPU", 0) < 2: + pytest.skip("DDP checkpoint test requires at least 2 GPUs in Ray cluster") + + yield + if started_cluster: + teardown_ray_cluster() + elif ray.is_initialized(): + ray.shutdown() + + +def _load_checkpoint(ckpt_path): + return torch.load(ckpt_path, map_location="cpu") + + +def _format_checkpoint(state): + formatted = {} + for module_name, module_state in state.items(): + formatted[module_name] = { + param_name: tensor.detach().cpu().tolist() for param_name, tensor in module_state.items() + } + return formatted + + +@pytest.mark.integration +@pytest.mark.gpu +def test_ddp_checkpoints_consistent(tmp_path, monkeypatch): + monkeypatch.setenv("RLIGHTNING_REMOTE_TRAIN", "1") + monkeypatch.setenv("RLIGHTNING_REMOTE_EVAL", "0") + assert InternalFlag.REMOTE_TRAIN is True + + policy_cfg = PolicyConfig( + type="DDPCheckpointPolicy", + rollout_mode="sync", + optim_cfg=Config.from_dict({"lr": 1e-2}), + weight_buffer=WeightBufferConfig(buffer_strategy="None"), + ) + train_config = TrainConfig(max_epochs=1, parallel="ddp") + + from rlightning.utils.builders import build_policy_group + + policy_group = build_policy_group( + policy_cls=policy_cfg.type, + policy_cfg=policy_cfg, + cluster_cfg=ClusterConfig(train_worker_num=2, eval_worker_num=0), + ) + + policy_group.init_train(train_config, env_meta=None) + + initial_paths = [] + initial_refs = [] + for idx, policy in enumerate(policy_group.train_list): + ckpt_path = tmp_path / f"initial_rank_{idx}_epoch_0.pth" + initial_paths.append(ckpt_path) + initial_refs.append(policy.save_checkpoint.remote(ckpt_path)) + + import ray + + ray.get(initial_refs) + + initial_states = [_load_checkpoint(ckpt_path) for ckpt_path in initial_paths] + # print("initial checkpoints:", [_format_checkpoint(state) for state in initial_states]) + + data = { + "x": torch.randn(16, 4), + "y": torch.randn(16, 2), + } + policy_group.update_dataset([data, data]) + + for _ in range(3): + policy_group.train() + + save_paths = [] + save_refs = [] + for idx, policy in enumerate(policy_group.train_list): + ckpt_path = tmp_path / f"rank_{idx}.pth" + save_paths.append(ckpt_path) + save_refs.append(policy.save_checkpoint.remote(ckpt_path)) + + ray.get(save_refs) + + states = [_load_checkpoint(ckpt_path) for ckpt_path in save_paths] + # print("after 3 steps checkpoints:", [_format_checkpoint(state) for state in states]) + assert states[0].keys() == states[1].keys() + for module_name in states[0]: + module_state_0 = states[0][module_name] + module_state_1 = states[1][module_name] + assert module_state_0.keys() == module_state_1.keys() + for param_name in module_state_0: + torch.testing.assert_close(module_state_0[param_name], module_state_1[param_name]) + for module_name in states[0]: + module_state_0 = states[0][module_name] + module_state_1 = states[1][module_name] + assert module_state_0.keys() == module_state_1.keys() + for param_name in module_state_0: + torch.testing.assert_close(module_state_0[param_name], module_state_1[param_name]) diff --git a/tests/integration/policy/test_policy_group.py b/tests/integration/policy/test_policy_group.py new file mode 100644 index 0000000..32fbb6f --- /dev/null +++ b/tests/integration/policy/test_policy_group.py @@ -0,0 +1,288 @@ +import os +from typing import Tuple + +import pytest +from gymnasium import spaces as gym_spaces +from tqdm import tqdm + +from rlightning.types import EnvMeta +from rlightning.utils.builders import ( + build_data_buffer, + build_env_group, + build_policy_group, +) +from rlightning.utils.config import ( + BufferConfig, + PolicyConfig, + TrainConfig, + WeightBufferConfig, + ClusterConfig, +) +from rlightning.utils.ray import resolve_object +from rlightning.utils.utils import InternalFlag +from tests.test_utils import _make_env_config, require_env_backend + +pytestmark = pytest.mark.integration + + +@pytest.fixture(scope="module", autouse=True) +def _auto_ray_cluster(): + # Ensure a Ray cluster is available for this module's tests + from tests.test_utils import setup_ray_cluster, teardown_ray_cluster + + setup_ray_cluster(num_cpus=4, num_gpus=4, local_mode=False) + yield + teardown_ray_cluster() + + +def _make_env_group(env_backend, task_name, as_remote): + require_env_backend(env_backend) + num_workers = 1 + num_envs = 1 + env_cfgs = [_make_env_config(env_backend, task_name, num_workers, num_envs)] + + os.environ["RLIGHTNING_REMOTE_ENV"] = "1" if as_remote else "0" + assert InternalFlag.REMOTE_ENV == as_remote + return build_env_group(env_cfgs=env_cfgs) + + +def _make_policy_config(): + return PolicyConfig( + type="SimplePPOPolicy", + weight_buffer=WeightBufferConfig( + type="WeightBuffer", + buffer_strategy="Double", + ), + rollout_mode="sync", + optim_cfg=WeightBufferConfig( + lr=1e-4, + ), + ) + + +def _make_train_config(): + return TrainConfig( + max_epochs=10, + lr=1e-3, + max_rollout_steps=10, + gamma=0.99, + ppo_epochs=1, + batch_size=10, + minibatch_size=2, + clip_ratio=0.2, + entropy_coef=1e-3, + value_coef=0.1, + max_grad_norm=0.5, + epochs=2, + ) + + +def _make_buffer_config(capacity: int = 100) -> BufferConfig: + return BufferConfig.from_dict( + { + "type": "RolloutBuffer", + "capacity": capacity, + "storage": { + "type": "unified", + "device": "cpu", + }, + "sampler": { + "type": "uniform", + }, + } + ) + + +def _make_buffer(buffer_cfg, env_meta_list=None, as_remote=False): + from tests.test_utils import _episode_postprocess_fn + + os.environ["RLIGHTNING_REMOTE_STORAGE"] = "1" if as_remote else "0" + assert InternalFlag.REMOTE_STORAGE == as_remote + + buffer = build_data_buffer( + buffer_cls=buffer_cfg.type, + buffer_cfg=buffer_cfg, + # obs_preprocessor=obs_preprocessor, + postprocess_fn=_episode_postprocess_fn, + ) + buffer.init(env_meta_list) + + return buffer + + +def _make_policy_group(policy_cfg, train_as_remote, eval_as_remote): + os.environ["RLIGHTNING_REMOTE_TRAIN"] = "1" if train_as_remote else "0" + os.environ["RLIGHTNING_REMOTE_EVAL"] = "1" if eval_as_remote else "0" + + return build_policy_group( + policy_cls=policy_cfg.type, + policy_cfg=policy_cfg, + cluster_cfg=ClusterConfig(), + ) + + +@pytest.mark.parametrize( + "env_backend,task_name", + [ + ["ale", "ALE/Alien-v5"], + ], +) +@pytest.mark.parametrize("eval_mode", [False, True]) +@pytest.mark.parametrize("env_as_remote", [True, False]) +@pytest.mark.parametrize("policy_as_remote", [True, False]) +def test_policy_group_with_simple_ppo( + env_backend, + task_name: str, + eval_mode: bool, + env_as_remote: bool, + policy_as_remote: bool, +): + env_group = _make_env_group(env_backend, task_name, env_as_remote) + + policy_cfg = _make_policy_config() + train_config = _make_train_config() + + policy_group = _make_policy_group( + policy_cfg, train_as_remote=policy_as_remote, eval_as_remote=policy_as_remote + ) + + env_meta = resolve_object(env_group.init()[0]) + + assert isinstance(env_meta, EnvMeta), type(env_meta) + + assert isinstance(env_meta.observation_space, gym_spaces.Space), type( + env_meta.observation_space + ) + + assert isinstance(env_meta.action_space, gym_spaces.Space), type(env_meta.action_space) + + if eval_mode: + episode_meta = policy_group.init_eval(None, env_meta) + else: + policy_group.init_train(train_config, env_meta) + + +@pytest.mark.parametrize( + "env_backend,task_name", + [ + ["ale", "ALE/Adventure-v5"], + ], +) +@pytest.mark.parametrize("env_as_remote", [True, False]) +@pytest.mark.parametrize("policy_as_remote", [True, False]) +def test_simple_ppo_evaluation( + env_backend: str, + task_name: str, + env_as_remote: bool, + policy_as_remote: bool, +): + """Test the functionality of rollouts + + The coverage indluces: + + 1. Support RGB input environments only + 2. The correctness of V-trace has not been validated yet + + Args: + env_backend (str): Registered environment backend + task_name (str): Registered environment name + """ + + max_rollout_steps: int = 50 + + env_group = _make_env_group(env_backend, task_name, env_as_remote) + + policy_cfg = _make_policy_config() + # train_config = _make_train_config() + + policy_group = _make_policy_group( + policy_cfg, train_as_remote=policy_as_remote, eval_as_remote=policy_as_remote + ) + + env_meta = resolve_object(env_group.init()[0]) + + assert isinstance(env_meta.observation_space, gym_spaces.Space), type( + env_meta.observation_space + ) + + assert isinstance(env_meta.action_space, gym_spaces.Space), type(env_meta.action_space) + + policy_group.init_eval(env_meta=env_meta) + + batched_env_ret, _ = env_group.reset(seed=0) + + for _ in tqdm(range(max_rollout_steps), disable=True): + batched_policy_resp = policy_group.rollout_batch(batched_env_ret) + batched_env_ret, _ = env_group.step(batched_policy_resp) + + +@pytest.mark.parametrize( + "env_backend,task_name", + [ + ["ale", "ALE/Adventure-v5"], + ], +) +@pytest.mark.parametrize("env_as_remote", [True, False]) +@pytest.mark.parametrize("buffer_as_remote", [True, False]) +@pytest.mark.parametrize("policy_as_remote", [True, False]) +def test_simple_ppo_training( + env_backend: str, + task_name: str, + env_as_remote: bool, + buffer_as_remote: bool, + policy_as_remote: bool, +): + """Test the functionality of SimplePPOPolicy + + The coverage includes: + + 1. Support RGB input environments only + 2. The correctness of V-trace has not been validated yet + + Args: + env_backend (str): Registered environment backend + task_name (str): Task name, or registered environment name + """ + + batch_size: int = 10 + + env_group = _make_env_group(env_backend, task_name, env_as_remote) + env_meta_list = resolve_object(env_group.init()) + env_meta = env_meta_list[0] + + policy_cfg = _make_policy_config() + train_config = _make_train_config() + policy_cfg.train_config = train_config + + # create buffer here + buffer_cfg = _make_buffer_config(capacity=100) + buffer = _make_buffer(buffer_cfg, env_meta_list, as_remote=buffer_as_remote) + + policy_group = _make_policy_group( + policy_cfg, train_as_remote=policy_as_remote, eval_as_remote=policy_as_remote + ) + + assert isinstance(env_meta.observation_space, gym_spaces.Space), type( + env_meta.observation_space + ) + assert isinstance(env_meta.action_space, gym_spaces.Space), type(env_meta.action_space) + + policy_group.init_train(train_config, env_meta) + policy_group.init_eval(None, env_meta) + + batched_env_ret, _ = env_group.reset(seed=0) + + for _ in range(train_config.epochs): + for _ in tqdm(range(train_config.max_rollout_steps), disable=True): + batched_policy_resp = policy_group.rollout_batch(batched_env_ret) + buffer.add_batched_transition(batched_env_ret, batched_policy_resp) + batched_env_ret, _ = env_group.step(batched_policy_resp) + + batched_policy_resp = policy_group.rollout_batch(batched_env_ret) + buffer.add_batched_transition(batched_env_ret, batched_policy_resp) + buffer.truncate_episodes(batched_policy_resp.ids()) + + data = buffer.sample(batch_size=batch_size) + policy_group.update_dataset(data) + info = policy_group.train() + print(info) diff --git a/tests/integration/utils/__init__.py b/tests/integration/utils/__init__.py new file mode 100644 index 0000000..b23639a --- /dev/null +++ b/tests/integration/utils/__init__.py @@ -0,0 +1 @@ +"""Utility integration tests.""" diff --git a/tests/integration/utils/test_remote_cls.py b/tests/integration/utils/test_remote_cls.py new file mode 100644 index 0000000..9f25e9a --- /dev/null +++ b/tests/integration/utils/test_remote_cls.py @@ -0,0 +1,47 @@ +import logging +import os + +import pytest +import ray +import torch + +from tests.test_utils import init_test_ray +from tests.tests_utils.remote_class_helper import Custom + +pytestmark = [pytest.mark.integration, pytest.mark.gpu] + + +@pytest.fixture(autouse=True) +def _test_ray_runtime(): + if not torch.cuda.is_available(): + pytest.skip("CUDA not available") + + python_path = os.environ.get("PYTHONPATH", "") + runtime_env = { + "env_vars": { + "RAY_DEBUG": "1", + "PYTHONPATH": python_path, + }, + } + try: + init_test_ray(num_gpus=1, local_mode=False, runtime_env=runtime_env) + except Exception as exc: + logging.error(f"Failed to initialize Ray for test: {exc}") + pytest.fail(f"Failed to initialize Ray for test: {exc}") + + logging.info("Connected to test Ray cluster.") + + yield + + if ray.is_initialized(): + ray.shutdown() + + +def test_remote(): + remote_cls = Custom.as_remote() + + instance = remote_cls.options(num_gpus=0.5).remote() + + res = ray.get(instance._get_gpu_ids.remote()) + print(".... gpu ids:", res) + assert isinstance(res, list) and len(res) == 1 diff --git a/tests/integration/weights/__init__.py b/tests/integration/weights/__init__.py new file mode 100644 index 0000000..e910be8 --- /dev/null +++ b/tests/integration/weights/__init__.py @@ -0,0 +1 @@ +"""Weight-buffer integration tests.""" diff --git a/tests/integration/weights/test_weight_buffer.py b/tests/integration/weights/test_weight_buffer.py new file mode 100644 index 0000000..4f9eafd --- /dev/null +++ b/tests/integration/weights/test_weight_buffer.py @@ -0,0 +1,268 @@ +import os + +import numpy as np +import pytest +import ray +import torch +from gymnasium import spaces as gym_spaces + +from rlightning.types import EnvMeta +from rlightning.utils.builders import build_policy_group +from rlightning.utils.config import Config, PolicyConfig, TrainConfig, WeightBufferConfig, ClusterConfig +from rlightning.weights.weight_buffer import WeightBuffer + +pytestmark = [pytest.mark.integration, pytest.mark.gpu] + +POLICY_AS_REMOTE = True + + +@pytest.fixture(scope="module", autouse=True) +def _auto_ray_cluster(): + # Ensure a Ray cluster is available for this module's tests + from tests.test_utils import setup_ray_cluster, teardown_ray_cluster + + if not torch.cuda.is_available(): + pytest.skip("CUDA not available") + + os.environ["RLIGHTNING_REMOTE_TRAIN"] = "1" if POLICY_AS_REMOTE else "0" + os.environ["RLIGHTNING_REMOTE_EVAL"] = "1" if POLICY_AS_REMOTE else "0" + + setup_ray_cluster(num_cpus=4, num_gpus=4, local_mode=False) + yield + teardown_ray_cluster() + + +def mock_weight_buffer_config(buffer_strategy: str = "Double"): + if buffer_strategy == "Double": + return WeightBufferConfig( + type="WeightBuffer", + buffer_strategy=buffer_strategy, + ) + elif buffer_strategy == "Shared": + return WeightBufferConfig( + type="CPUWeightBuffer", + buffer_strategy=buffer_strategy, + ) + + +def mock_env_meta(): + """Build a minimal env_meta compatible with SimplePPOPolicy.""" + action_space = gym_spaces.Discrete(5) + return EnvMeta( + env_id=None, + action_space=action_space, + ) + + +def mock_policy_config(buffer_strategy: str = "Double"): + return PolicyConfig( + type="SimplePPOPolicy", + rollout_mode="sync", + weight_buffer=mock_weight_buffer_config(buffer_strategy), + optim_cfg=Config(lr=0.001), + train_config=TrainConfig(max_epochs=1), + ) + + +def mock_policy_group(policy_cfg): + return build_policy_group( + policy_cls=policy_cfg.type, + policy_cfg=policy_cfg, + cluster_cfg=ClusterConfig(), + ) + + +def _get_policy_state_dict(policy, policy_as_remote): + model_parameters = ( + ray.get(policy.get_trainable_parameters.remote()) + if policy_as_remote + else policy.get_trainable_parameters() + ) + return model_parameters + + +def _resolve_weight_buffer(eval_policy): + """Resolve the actual weight buffer object/actor from a policy.""" + if isinstance(eval_policy, ray.actor.ActorHandle): + return ray.get(eval_policy.get_weight_buffer.remote()) + return eval_policy.get_weight_buffer() + + +def _get_weight_buffer_state_dict(weight_buffer): + """Helper to get state dict from weight buffer.""" + if isinstance(weight_buffer, ray.actor.ActorHandle): + return ray.get(weight_buffer.get_state_dict.remote()) + return weight_buffer.get_state_dict() + + +def _add_state_dict_to_buffer(weight_buffer, state_dict): + """Helper to add a state dict to weight buffer.""" + if isinstance(weight_buffer, ray.actor.ActorHandle): + ray.get(weight_buffer.add_state_dict.remote(state_dict)) + else: + weight_buffer.add_state_dict(state_dict) + + +@pytest.mark.parametrize("policy_as_remote", [POLICY_AS_REMOTE]) +@pytest.mark.parametrize("buffer_strategy", ["Double", "Shared"]) +def test_init_weight_buffer_with_simple_ppo( + policy_as_remote: bool, + buffer_strategy: str, +): + """Test weight buffer initialization with SimplePPO policy""" + policy_cfg = mock_policy_config(buffer_strategy=buffer_strategy) + + # Initialize weight buffer for eval policy + policy_group = mock_policy_group(policy_cfg) + + eval_policy = policy_group.eval_list[0] + weight_buffer = _resolve_weight_buffer(eval_policy) + + # Verify weight buffer is initialized + assert weight_buffer is not None + if buffer_strategy == "Shared": + assert isinstance(weight_buffer, ray.actor.ActorHandle) + elif buffer_strategy == "Double": + assert isinstance(weight_buffer, WeightBuffer) + + +@pytest.mark.parametrize("policy_as_remote", [POLICY_AS_REMOTE]) +@pytest.mark.parametrize("buffer_strategy", ["Double", "Shared"]) +def test_send_weights_with_simple_ppo( + policy_as_remote: bool, + buffer_strategy: str, +): + """Test sending weights with SimplePPO policy""" + policy_cfg = mock_policy_config(buffer_strategy=buffer_strategy) + + # Initialize weight buffer for eval policy + policy_group = mock_policy_group(policy_cfg) + + # Initialize eval and train policies + env_meta = mock_env_meta() + policy_group.init_eval(env_meta=env_meta) + policy_group.init_train(policy_cfg.train_config, env_meta=env_meta) + + train_policy = policy_group.train_list[0] + eval_policy = policy_group.eval_list[0] + + # send weights to eval policy + if policy_as_remote: + if buffer_strategy == "Double": + ray.get(train_policy.send_weights.remote([eval_policy])) + weight_buffer = _resolve_weight_buffer(eval_policy) + elif buffer_strategy == "Shared": + weight_buffer = next(iter(policy_group.shared_weight_buffer_map.values()))[ + "shared_weight_buffer" + ] + ray.get(train_policy.send_weights.remote([], shared_weight_buffer=weight_buffer)) + + # Verify weights were added + retrieved_dict = _get_weight_buffer_state_dict(weight_buffer) + assert retrieved_dict != {}, "Initial state dict is empty" + + # Create a state dict with model parameters + initial_state_dict = _get_policy_state_dict(train_policy, policy_as_remote) + + # Verify the state dict contains expected keys + for module_name, module_state_dict in retrieved_dict.items(): + print(module_name) + for name in module_state_dict.keys(): + assert name in initial_state_dict[module_name] + assert module_state_dict[name].shape == initial_state_dict[module_name][name].shape + + +@pytest.mark.parametrize("policy_as_remote", [POLICY_AS_REMOTE]) +@pytest.mark.parametrize("buffer_strategy", ["Double", "Shared"]) +def test_update_weights_with_simple_ppo( + policy_as_remote: bool, + buffer_strategy: str, +): + """Test updating weights with SimplePPO policy""" + policy_cfg = mock_policy_config(buffer_strategy=buffer_strategy) + + # Initialize weight buffer for eval policy + policy_group = mock_policy_group(policy_cfg) + + # Initialize eval and train policies + env_meta = mock_env_meta() + policy_group.init_eval(env_meta=env_meta) + policy_group.init_train(policy_cfg.train_config, env_meta=env_meta) + + train_policy = policy_group.train_list[0] + eval_policy = policy_group.eval_list[0] + + # send weights to eval policy + if policy_as_remote: + if buffer_strategy == "Double": + ray.get(train_policy.send_weights.remote([eval_policy])) + weight_buffer = _resolve_weight_buffer(eval_policy) + elif buffer_strategy == "Shared": + weight_buffer = next(iter(policy_group.shared_weight_buffer_map.values()))[ + "shared_weight_buffer" + ] + ray.get(train_policy.send_weights.remote([], shared_weight_buffer=weight_buffer)) + + # Get initial weights from train policy + initial_state_dict = _get_policy_state_dict(train_policy, policy_as_remote) + + ray.get(eval_policy.update_weights_from_buffer.remote()) + eval_state_dict = _get_policy_state_dict(eval_policy, policy_as_remote) + + # Verify eval policy weights were updated from buffer + for module_name, module_state_dict in eval_state_dict.items(): + print(module_name) + for name in module_state_dict.keys(): + assert name in initial_state_dict[module_name] + assert torch.allclose( + eval_state_dict[module_name][name], initial_state_dict[module_name][name] + ) + + +@pytest.mark.parametrize("buffer_strategy", ["Double", "Shared"]) +def test_weight_buffer_clear_and_reinit(buffer_strategy: str): + """Test clearing and reinitializing weight buffer""" + from rlightning.weights.utils import build_weight_buffer + + weight_buffer_cfg = mock_weight_buffer_config(buffer_strategy) + node_id = ray.get_runtime_context().get_node_id() + weight_buffer = build_weight_buffer(weight_buffer_cfg.type, weight_buffer_cfg, node_id=node_id) + + # First state dict + state_dict1 = { + "layer1": { + "weight": ( + torch.randn(10, 10) if buffer_strategy == "Double" else np.random.randn(10, 10) + ) + } + } + _add_state_dict_to_buffer(weight_buffer, state_dict1) + + # Check readiness + if isinstance(weight_buffer, ray.actor.ActorHandle): + is_ready = ray.get(weight_buffer.is_ready.remote()) + else: + is_ready = weight_buffer.is_ready + assert is_ready + + # Clear buffer + if isinstance(weight_buffer, ray.actor.ActorHandle): + ray.get(weight_buffer.clear.remote()) + is_ready = ray.get(weight_buffer.is_ready.remote()) + else: + weight_buffer.clear() + is_ready = weight_buffer.is_ready() + assert not is_ready + + # Add new state dict + state_dict2 = { + "layer2": { + "weight": torch.randn(5, 5) if buffer_strategy == "Double" else np.random.randn(5, 5) + } + } + _add_state_dict_to_buffer(weight_buffer, state_dict2) + + # Verify new state dict is present + retrieved = _get_weight_buffer_state_dict(weight_buffer) + assert "layer2" in retrieved + assert "layer1" not in retrieved diff --git a/tests/test_utils.py b/tests/test_utils.py new file mode 100644 index 0000000..2ef32d0 --- /dev/null +++ b/tests/test_utils.py @@ -0,0 +1,453 @@ +"""Shared helper utilities for tests.""" + +from __future__ import annotations + +import os +import random +import time +import uuid +from unittest.mock import MagicMock + +import numpy as np +import pytest +import torch +from omegaconf import OmegaConf + +from rlightning.utils.config import EnvConfig + +try: + import gymnasium as gym +except Exception: # pragma: no cover - fallback only used in degraded envs. + from types import SimpleNamespace + + class _DummyEnv: + def reset(self, seed=None, options=None): + return None + + class _DummySpace: + pass + + class _DummyBox(_DummySpace): + def __init__(self, low, high, shape, dtype=None): + self.low = low + self.high = high + self.shape = shape + self.dtype = dtype + + class _DummyDict(dict): + pass + + gym = SimpleNamespace(Env=_DummyEnv, spaces=SimpleNamespace(Box=_DummyBox, Dict=_DummyDict)) + +try: + import ray + + RAY_AVAILABLE = True +except Exception: + RAY_AVAILABLE = False + +try: + import vllm # noqa: F401 + + VLLM_AVAILABLE = True +except Exception: + VLLM_AVAILABLE = False + +try: + import mani_skill # noqa: F401 + + MANISKILL_AVAILABLE = True +except Exception: + MANISKILL_AVAILABLE = False + + +class MockEnv(gym.Env): + def __init__(self, config=None, env_id="mock_env", **kwargs): + super().__init__() + self.config = config or {} + self.env_id = env_id + self.observation_space = gym.spaces.Dict( + { + "rgb": gym.spaces.Box(0, 255, (64, 64, 3), dtype=np.uint8), + "state": gym.spaces.Box(-np.inf, np.inf, (10,), dtype=np.float32), + } + ) + self.action_space = gym.spaces.Box(-1, 1, (4,), dtype=np.float32) + self.step_count = 0 + self.max_steps = 100 + + def reset(self, seed=None, options=None): + super().reset(seed=seed) + self.step_count = 0 + obs = { + "rgb": np.random.randint(0, 255, (64, 64, 3), dtype=np.uint8), + "state": np.random.randn(10).astype(np.float32), + } + return obs, {} + + def step(self, action): + self.step_count += 1 + obs = { + "rgb": np.random.randint(0, 255, (64, 64, 3), dtype=np.uint8), + "state": np.random.randn(10).astype(np.float32), + } + reward = float(np.random.randn()) + terminated = self.step_count >= self.max_steps + truncated = False + info = {"truncated": truncated} + return obs, reward, terminated, truncated, info + + def init(self, *args, **kwargs): + return self.observation_space, self.action_space + + +class MockEnvGroup: + def __init__(self, env): + self.env = env + self._observation_space = env.observation_space + self._action_space = env.action_space + + def init(self): + return self._observation_space, self._action_space + + def reset(self, seed=0): + obs, info = self.env.reset(seed=seed) + return MagicMock(obs=[obs], info=[info]), [] + + def step(self, responses): + actions = [r.action for r in responses] if hasattr(responses[0], "action") else responses + obs, reward, terminated, truncated, info = self.env.step(actions[0]) + return ( + MagicMock(obs=[obs], reward=[reward], done=[terminated], truncated=[truncated], info=[info]), + [], + ) + + def observation_space(self): + return self._observation_space + + def action_space(self): + return self._action_space + + +class MockPolicyGroup: + def __init__(self): + self.task = None + self.train_config = {} + + def init_eval(self, eval_config, env_meta): + return MagicMock() + + def init_train(self, train_config, env_meta): + pass + + def rollout_batch(self, requests): + actions = torch.randn(len(requests.obs), 4) + return MagicMock( + action=actions, + value=torch.randn(len(requests.obs)), + log_prob=torch.randn(len(requests.obs)), + ids=lambda: [str(uuid.uuid4()) for _ in range(len(requests.obs))], + ) + + def update_dataset(self, buffer, batch_size): + pass + + def train(self): + pass + + def send_weights_to_eval(self): + pass + + def notify_update_weights(self): + pass + + +class MockBuffer: + def __init__(self, config): + self.config = config + self.data = [] + + def add(self, obs, act): + self.data.append((obs, act)) + + def sample(self, batch_size, device="cpu"): + if not self.data: + return None, None + indices = np.random.choice(len(self.data), min(batch_size, len(self.data)), replace=False) + obs_batch = [self.data[i][0] for i in indices] + act_batch = [self.data[i][1] for i in indices] + return obs_batch, act_batch + + def init(self, episode_meta): + pass + + def add_transitions_async(self, transitions, truncations=None): + if hasattr(transitions, "obs"): + for i in range(len(transitions.obs)): + self.add(transitions.obs[i], getattr(transitions, "action", [0])[i]) + + def size(self): + return len(self.data) + + def clear(self): + self.data.clear() + + def truncate_episodes(self, episode_ids): + pass + + +def create_tiny_config(): + return OmegaConf.create( + { + "env": { + "envs": [{"name": "mock_env", "num_envs": 1}], + "num_envs": 1, + }, + "policy": { + "type": "mock_policy", + "train_config": {}, + }, + "buffer": { + "type": "mock_buffer", + }, + "train": { + "max_epochs": 2, + "max_rollout_steps": 10, + "batch_size": 4, + }, + "mode": "sync", + "train_worker_num": 1, + "eval_worker_num": 1, + "logging": { + "type": "console", + }, + } + ) + + +def create_sample_obs(): + return { + "rgb": torch.randint(0, 255, (1, 64, 64, 3), dtype=torch.uint8), + "state": torch.randn(1, 10), + } + + +def create_sample_episode(): + t = 10 + return { + "obs": [torch.randn(64, 64, 3) for _ in range(t)], + "action": [torch.randn(4) for _ in range(t)], + "reward": [torch.tensor(float(np.random.randn())) for _ in range(t)], + "value": [torch.tensor(float(np.random.randn())) for _ in range(t)], + "next/value": [torch.tensor(float(np.random.randn())) for _ in range(t)], + "done": [torch.tensor(False) for _ in range(t - 1)] + [torch.tensor(True)], + "truncated": [torch.tensor(False) for _ in range(t)], + } + + +_ORIGINAL_CUDA_VISIBLE_DEVICES = os.environ.get("CUDA_VISIBLE_DEVICES", None) +_ORIG_CUDA = None + + +def get_test_ray_address(default="local"): + return os.environ.get("RLIGHTNING_TEST_RAY_ADDRESS", default) + + +def init_test_ray( + num_cpus=1, + num_gpus=0, + local_mode=True, + log_to_driver=False, + runtime_env=None, + default_address="local", +): + global _ORIG_CUDA + + if not RAY_AVAILABLE: + raise RuntimeError("Ray not available") + + if _ORIG_CUDA is None: + _ORIG_CUDA = os.environ.get("CUDA_VISIBLE_DEVICES", None) + + if ray.is_initialized(): + ray.shutdown() + time.sleep(0.2) + + ray_address = get_test_ray_address(default=default_address) + if ray_address == "local": + ray.init( + address="local", + num_cpus=num_cpus, + num_gpus=num_gpus, + ignore_reinit_error=True, + include_dashboard=False, + log_to_driver=log_to_driver, + local_mode=local_mode, + runtime_env=runtime_env, + ) + else: + ray.init( + address=ray_address, + ignore_reinit_error=True, + log_to_driver=log_to_driver, + runtime_env=runtime_env, + ) + + return ray_address + + +def setup_ray_cluster(num_cpus=1, num_gpus=0, local_mode=True, log_to_driver=False): + return init_test_ray( + num_cpus=num_cpus, + num_gpus=num_gpus, + local_mode=local_mode, + log_to_driver=log_to_driver, + default_address="local", + ) + + +def teardown_ray_cluster(): + global _ORIG_CUDA + + if RAY_AVAILABLE and ray.is_initialized(): + ray.shutdown() + time.sleep(0.2) + + if _ORIG_CUDA is None: + os.environ.pop("CUDA_VISIBLE_DEVICES", None) + else: + os.environ["CUDA_VISIBLE_DEVICES"] = _ORIG_CUDA + + +def check_dependency_availability(): + return { + "ray": RAY_AVAILABLE, + "vllm": VLLM_AVAILABLE, + "maniskill": MANISKILL_AVAILABLE, + "gpu": torch.cuda.is_available(), + } + + +def create_tmp_artifacts_dir(tmp_path): + artifacts_dir = tmp_path / "outputs" / "test_artifacts" + artifacts_dir.mkdir(parents=True, exist_ok=True) + return artifacts_dir + + +def set_deterministic_seeds(seed=42): + np.random.seed(seed) + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + random.seed(seed) + + +def _compute_gae(rewards, values, dones, gamma, gae_lambda): + advantages = torch.zeros_like(rewards) + lastgaelam = torch.zeros_like(rewards[0]) + + for t in reversed(range(rewards.shape[0])): + next_not_done = 1.0 - dones[t].float() + current_value = values[t] + next_value = values[t + 1] + delta = rewards[t] + gamma * next_value * next_not_done - current_value + lastgaelam = delta + gamma * gae_lambda * next_not_done * lastgaelam + advantages[t] = lastgaelam + + returns = advantages + values[:-1] + return advantages, returns + + +def _episode_postprocess_fn(raw_episode): + episode = {} + for key, value in raw_episode.items(): + if "info" in key: + continue + if isinstance(value[0], torch.Tensor): + episode[key] = torch.stack(value).squeeze() + else: + episode[key] = torch.tensor(value) + + episode["last_reward"] = episode["last_reward"][1:] + dones = torch.logical_or(episode["last_terminated"], episode["last_truncated"]) + episode["done"] = dones[1:] + episode["action"] = episode["action"][:-1] + episode["last_terminated"] = episode["last_terminated"][1:] + episode["last_truncated"] = episode["last_truncated"][1:] + episode["log_prob"] = episode["log_prob"][:-1] + episode["entropy"] = episode["entropy"][:-1] + episode["observation"] = episode["observation"][:-1] + + rewards = episode["last_reward"] + values = episode["value"] + episode["value"] = episode["value"][:-1] + dones = episode["done"] + + advantages, returns = _compute_gae(rewards, values, dones, 0.8, 0.9) + advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8) + + episode["advantages"] = advantages + episode["returns"] = returns + return episode + + +def _make_env_config(env_backend, task_name, num_workers, num_envs=1): + return EnvConfig( + backend=env_backend, + name=f"{env_backend}/{task_name}", + task=task_name, + max_episode_steps=20, + num_workers=num_workers, + num_envs=num_envs, + ) + + +class PerformanceTimer: + def __init__(self): + self.start_time = None + self.end_time = None + + def start(self): + self.start_time = time.time() + + def stop(self): + self.end_time = time.time() + return self.elapsed + + @property + def elapsed(self): + if self.start_time is None: + return 0 + end = self.end_time or time.time() + return end - self.start_time + + +def should_skip_gpu_test(): + return not torch.cuda.is_available() + + +def should_skip_ray_test(): + return not RAY_AVAILABLE + + +def should_skip_vllm_test(): + return not VLLM_AVAILABLE + + +def should_skip_maniskill_test(): + return not MANISKILL_AVAILABLE + + +def require_env_backend(backend: str) -> None: + if backend == "mujoco": + pytest.importorskip("mujoco") + elif backend == "ale": + pytest.importorskip("ale_py") + elif backend == "isaac_manager_based": + pytest.importorskip("isaaclab") + + +@pytest.fixture +def performance_timer(): + return PerformanceTimer() diff --git a/tests/tests_utils/__init__.py b/tests/tests_utils/__init__.py new file mode 100644 index 0000000..efff8eb --- /dev/null +++ b/tests/tests_utils/__init__.py @@ -0,0 +1 @@ +"""Integration test helper modules.""" diff --git a/tests/tests_utils/ddp_checkpoint_policy.py b/tests/tests_utils/ddp_checkpoint_policy.py new file mode 100644 index 0000000..9d671bd --- /dev/null +++ b/tests/tests_utils/ddp_checkpoint_policy.py @@ -0,0 +1,51 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.nn.parallel import DistributedDataParallel as DDP + +from rlightning.policy.base_policy import BasePolicy +from rlightning.utils.registry import POLICIES + + +@POLICIES.register("DDPCheckpointPolicy") +class DDPCheckpointPolicy(BasePolicy): + def construct_network(self, env_meta=None, *args, **kwargs): + self.linear = nn.Linear(4, 2, bias=True) + nn.init.constant_(self.linear.weight, 0.1) + nn.init.constant_(self.linear.bias, 0.0) + if torch.cuda.is_available(): + self.linear.cuda() + + def setup_optimizer(self, optim_cfg): + lr = getattr(optim_cfg, "lr", 1e-2) + self.optimizer = torch.optim.SGD(self.linear.parameters(), lr=lr) + + def update_dataset(self, data): + self._dataset = data + + def train(self): + x = self._dataset["x"].cuda() + y = self._dataset["y"].cuda() + pred = self.linear(x) + loss = F.mse_loss(pred, y) + self.optimizer.zero_grad() + loss.backward() + self.optimizer.step() + return {"loss": float(loss.detach().cpu())} + + def rollout_step(self, env_ret): + raise NotImplementedError("DDP checkpoint test policy does not support rollout.") + + def postprocess(self, data): + raise NotImplementedError("DDP checkpoint test policy does not support postprocess.") + + def get_trainable_parameters(self): + state_dict = {} + for name, model in self.model_list: + module = model.module if isinstance(model, DDP) else model + state_dict[name] = module.state_dict() + return state_dict + + def load_state_dict(self, state_dict, *args, **kwargs): + for name, model in self.model_list: + model.load_state_dict(state_dict[name], strict=True) diff --git a/tests/tests_utils/envs/__init__.py b/tests/tests_utils/envs/__init__.py new file mode 100644 index 0000000..b53fc37 --- /dev/null +++ b/tests/tests_utils/envs/__init__.py @@ -0,0 +1 @@ +"""Environment-side integration helpers.""" diff --git a/tests/tests_utils/envs/utils_remote_env.py b/tests/tests_utils/envs/utils_remote_env.py new file mode 100644 index 0000000..6c365a2 --- /dev/null +++ b/tests/tests_utils/envs/utils_remote_env.py @@ -0,0 +1,165 @@ +import random +import time +from typing import Dict, List + +import numpy as np +import ray +import torch + +from rlightning.env import BaseEnv +from rlightning.env.env_server import RemoteEnvServer +from rlightning.env.remote_env.env_client import RemoteEnvClient +from rlightning.env.utils.utils import default_env_preprocess_fn +from rlightning.types import EnvRet, PolicyResponse, Processed_EnvRet_fields +from rlightning.utils.config import EnvConfig + + +class MockPiperEnv(BaseEnv): + def __init__(self, config, worker_index=0, preprocess_fn=default_env_preprocess_fn) -> None: + super().__init__(config, worker_index, preprocess_fn) + self.step_cnt = 0 + self.episode_cnt = 0 + self.episode_reward = 0 + self.total_steps = config.total_steps + + def reset(self): + obs = { + "state": np.zeros((7,)), + "camera_rgb_front": np.zeros((480, 640, 3), dtype=np.uint8), + "camera_rgb_wrist": np.zeros((480, 640, 3), dtype=np.uint8), + } + self.step_cnt = 0 + self.episode_cnt = 0 + self.episode_reward = 0 + return EnvRet(env_id=self.env_id, observation=obs, ts_env_sent_ns=time.time_ns()) + + def step(self, policy_resp: PolicyResponse) -> EnvRet: + _ = self._preprocess_fn(policy_resp) + obs = { + "state": np.zeros((7,)), + "camera_rgb_front": np.zeros((480, 640, 3), dtype=np.uint8), + "camera_rgb_wrist": np.zeros((480, 640, 3), dtype=np.uint8), + } + done = self.step_cnt == random.choice([2, 5, 8]) + truncated = self.step_cnt == random.choice([2, 5, 8]) + self.step_cnt += 1 + time.sleep(random.random() * 0.2) + return EnvRet( + env_id=policy_resp.env_id, + observation=obs, + last_reward=0.0, + last_terminated=done, + last_truncated=truncated, + info={}, + ts_env_sent_ns=time.time_ns(), + ) + + @classmethod + def episode_postprocess_fn(cls, episode_buffer: Dict) -> Dict: + data = {} + for k, v in episode_buffer.items(): + if "info" in k: + continue + if k == "observation": + state = np.array([elem["state"] for elem in v]) + data["state"] = state[:-1] + data["next_state"] = state[1:] + camera_rgb_front = np.array([elem["camera_rgb_front"] for elem in v]) + data["camera_rgb_front"] = camera_rgb_front[:-1] + data["next_camera_rgb_front"] = camera_rgb_front[1:] + camera_rgb_wrist = np.array([elem["camera_rgb_wrist"] for elem in v]) + data["camera_rgb_wrist"] = camera_rgb_wrist[:-1] + data["next_camera_rgb_wrist"] = camera_rgb_wrist[1:] + continue + + if isinstance(v[0], torch.Tensor): + v = torch.stack(v, dim=0) + else: + v = torch.tensor(v) + if k.startswith("last_"): + data[k[5:]] = v[1:] + else: + data[k] = v + + env_fields = set(EnvRet.fields() + Processed_EnvRet_fields) + policy_fields = [field for field in episode_buffer.keys() if field not in env_fields] + if k in policy_fields: + data[k] = v[:-1] + + return data + + def get_action_space(self): + return + + def get_observation_space(self): + return + + def is_finish(self): + return self.step_cnt >= self.total_steps + + +@ray.remote +class EnvClientWorker: + def __init__(self, address, port, total_steps=10): + config = EnvConfig( + name="MockPiper-v0", + task="MockPiper-v0", + backend="piper", + max_episode_steps=1000, + total_steps=total_steps, + ) + self.env = MockPiperEnv(config) + self.client = RemoteEnvClient(self.env, address, port) + + def run(self): + self.client.connect() + self.client.run() + time.sleep(random.choice([1, 4, 7])) + self.client.env.reset() + self.client.connect() + self.client.run() + + def is_alive(self): + return True + + +@ray.remote +class EnvServerWorker: + def __init__(self): + config = EnvConfig(name="test_env_server", backend="env_server", task="real_world", zmq_port="6366") + self.server = RemoteEnvServer(config) + self.server.init() + + def _mock_policy_response_list(self, env_ret_list: List[EnvRet]) -> List[PolicyResponse]: + policy_response_list = [] + for env_ret in env_ret_list: + action = np.random.uniform(-1, 1, size=(7,)) + policy_response_list.append(PolicyResponse(env_id=env_ret.env_id, action=action)) + return policy_response_list + + def get_address_port(self): + return self.server.get_address_port() + + def run(self, num_envs, num_steps, timeout=100): + server = self.server + env_ret_list = server.reset() + cnt = len(env_ret_list) + expect_cnts = (1 + num_steps) * num_envs * 2 + + while True: + policy_resp_list = self._mock_policy_response_list(env_ret_list) + server.step_async(policy_resp_list) + env_ret_list = server.collect_async() + cnt += len(env_ret_list) + + if cnt == expect_cnts: + policy_resp_list = self._mock_policy_response_list(env_ret_list) + server.step_async(policy_resp_list) + env_ret_list = server.collect_async(timeout=10) + return len(env_ret_list) == 0 + + def is_alive(self): + return True + + def close(self): + self.server.close() diff --git a/tests/tests_utils/mini_experiment.py b/tests/tests_utils/mini_experiment.py new file mode 100644 index 0000000..1f1f0be --- /dev/null +++ b/tests/tests_utils/mini_experiment.py @@ -0,0 +1,242 @@ +"""Test-only components for running tiny end-to-end training experiments.""" + +from __future__ import annotations + +import datetime +import os +from typing import Any, Dict + +import numpy as np +import torch +import torch.nn.functional as F +from gymnasium import spaces +from torch import nn +from torch import distributed as dist +from types import SimpleNamespace + +from rlightning.env.base_env import BaseEnv +from rlightning.policy.base_policy import BasePolicy +from rlightning.types import EnvRet, PolicyResponse +from rlightning.utils.distributed.comm_context import CommContext +from rlightning.utils.distributed.group_initializer import CommMode +from rlightning.utils.registry import ENVS, POLICIES +from rlightning.utils.ray.remote_class import RayActorMixin +from rlightning.utils.utils import to_device + + +if not torch.cuda.is_available(): + # BasePolicy currently calls `.cuda()` in rollout/postprocess hooks even on CPU-only + # machines. Patch the data containers in this test-only module so subprocess smoke + # tests can exercise the full engine stack without requiring GPUs. + EnvRet.cuda = lambda self, device="cuda": self # type: ignore[method-assign] + PolicyResponse.cuda = lambda self, device="cuda": self # type: ignore[method-assign] + RayActorMixin._get_gpu_ids = lambda self: [0] # type: ignore[method-assign] + + def _cpu_safe_init_distributed_env( + self, + world_size: int | None = None, + rank: int | None = None, + backend: str = "gloo", + dist_url: str = "env://", + timeout: int = 1800, + ) -> None: + dist.init_process_group( + backend=backend, + init_method=dist_url, + world_size=world_size, + rank=rank, + timeout=datetime.timedelta(0, timeout), + ) + dist.barrier() + + ranks = list(range(world_size)) + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + self._register_group(local_rank, world_size, dist.GroupMember.WORLD, ranks, CommMode.GLOBAL) + self._global_ranks[CommMode.GLOBAL] = rank + + CommContext.init_distributed_env = _cpu_safe_init_distributed_env # type: ignore[method-assign] + + +@ENVS.register("mini_counter") +class MiniCounterEnv(BaseEnv): + """A tiny deterministic env whose target action is encoded in the observation.""" + + def __init__(self, config, worker_index=0, preprocess_fn=None): + super().__init__(config=config, worker_index=worker_index, preprocess_fn=preprocess_fn) + self.observation_space = spaces.Box(low=-1.0, high=1.0, shape=(4,), dtype=np.float32) + self.action_space = spaces.Box(low=-1.0, high=1.0, shape=(2,), dtype=np.float32) + self.env = SimpleNamespace( + observation_space=self.observation_space, + action_space=self.action_space, + ) + self.horizon = int(getattr(config, "max_episode_steps", None) or 4) + self._step_count = 0 + self._episode_return = 0.0 + self._obj_set = "train" + + def _target(self) -> np.ndarray: + if self._obj_set == "test": + return np.array([-0.5, 0.25], dtype=np.float32) + return np.array([0.25, -0.25], dtype=np.float32) + + def _observation(self) -> np.ndarray: + target = self._target() + return np.array( + [ + self._step_count / max(self.horizon, 1), + float(target[0]), + float(target[1]), + 1.0, + ], + dtype=np.float32, + ) + + def reset(self, seed=None, options=None): + del seed + self._step_count = 0 + self._episode_return = 0.0 + self._obj_set = (options or {}).get("obj_set", "train") + return EnvRet(env_id=self.env_id, observation=self._observation()) + + def step(self, policy_resp): + action = self._preprocess_fn(policy_resp) if self._preprocess_fn is not None else policy_resp.action + if isinstance(action, torch.Tensor): + action = action.detach().cpu().numpy() + action = np.asarray(action, dtype=np.float32).reshape(-1) + + target = self._target() + reward = 1.0 - float(np.square(action[:2] - target).mean()) + + self._step_count += 1 + self._episode_return += reward + terminated = self._step_count >= self.horizon + + info = {} + if terminated: + info["episode_info"] = { + "episode_return": self._episode_return, + "success": float(reward > 0.8), + } + + return EnvRet( + env_id=self.env_id, + observation=self._observation(), + last_reward=reward, + last_terminated=terminated, + last_truncated=False, + info=info, + ) + + +class _MiniPolicyBase(BasePolicy): + """Shared training logic for tiny sync/async smoke-test policies.""" + + def construct_network(self, env_meta, *args, **kwargs): + del args, kwargs + obs_shape = env_meta.observation_space.shape + act_shape = env_meta.action_space.shape + obs_dim = int(np.prod(obs_shape)) + action_dim = int(np.prod(act_shape)) + + self.encoder = nn.Sequential( + nn.Linear(obs_dim, 32), + nn.Tanh(), + ) + self.actor_head = nn.Linear(32, action_dim) + self.critic_head = nn.Linear(32, 1) + self.to(self.device) + + def setup_optimizer(self, optim_cfg): + lr = float(getattr(optim_cfg, "lr", 5e-2)) + self.optimizer = torch.optim.Adam(self.parameters(), lr=lr) + + def _tensor_obs(self, observation: Any) -> torch.Tensor: + obs = torch.as_tensor(observation, dtype=torch.float32, device=self.device) + if obs.ndim == 1: + obs = obs.unsqueeze(0) + return obs.reshape(obs.shape[0], -1) + + def _policy_forward(self, observation: Any) -> tuple[torch.Tensor, torch.Tensor]: + obs = self._tensor_obs(observation) + hidden = self.encoder(obs) + action = torch.tanh(self.actor_head(hidden)) + value = self.critic_head(hidden).squeeze(-1) + return action, value + + def _rollout_step_impl(self, env_ret: EnvRet) -> PolicyResponse: + action, value = self._policy_forward(env_ret.observation) + log_prob = torch.zeros_like(value) + return PolicyResponse( + env_id=env_ret.env_id, + action=action.squeeze(0), + log_prob=log_prob.squeeze(0), + value=value.squeeze(0), + ) + + def postprocess(self, env_ret=None, policy_resp=None): + return { + "env_id": getattr(env_ret, "env_id", getattr(policy_resp, "env_id", None)), + } + + def update_dataset(self, data) -> None: + self.dataset = to_device(data, self.device) + + def train(self): + obs = self.dataset["observation"].float() + obs = obs.reshape(obs.shape[0], -1) + reward = self.dataset["reward"].float().reshape(-1) + target_action = obs[:, 1:3].detach() + + inner_updates = int(getattr(self.config.train_config, "inner_updates", 2)) + metrics: Dict[str, float] = {} + for _ in range(inner_updates): + hidden = self.encoder(obs) + predicted_action = torch.tanh(self.actor_head(hidden)) + predicted_value = self.critic_head(hidden).squeeze(-1) + + actor_loss = F.mse_loss(predicted_action, target_action) + value_loss = F.mse_loss(predicted_value, reward) + loss = actor_loss + 0.1 * value_loss + + self.optimizer.zero_grad() + loss.backward() + self.optimizer.step() + + metrics = { + "loss": float(loss.detach().cpu().item()), + "actor_loss": float(actor_loss.detach().cpu().item()), + "value_loss": float(value_loss.detach().cpu().item()), + } + + return metrics + + def get_trainable_parameters(self) -> Dict[str, Dict[str, torch.Tensor]]: + params: Dict[str, Dict[str, torch.Tensor]] = {} + for name, model in self.model_list: + module = model.module if hasattr(model, "module") else model + params[name] = module.state_dict() + return params + + def load_state_dict(self, state_dict, strict=True, assign=False): + del strict, assign + for name, model in self.model_list: + module = model.module if hasattr(model, "module") else model + module.load_state_dict(state_dict[name]) + + +@POLICIES.register("MiniSyncPolicy") +class MiniSyncPolicy(_MiniPolicyBase): + """Sync rollout variant for tiny smoke tests.""" + + def rollout_step(self, env_ret: EnvRet, **kwargs): + del kwargs + return self._rollout_step_impl(env_ret) + + +@POLICIES.register("MiniAsyncPolicy") +class MiniAsyncPolicy(_MiniPolicyBase): + """Async rollout variant for tiny smoke tests.""" + + async def rollout_step(self, env_ret: EnvRet, **kwargs): + del kwargs + return self._rollout_step_impl(env_ret) diff --git a/tests/tests_utils/mini_experiment_conf/asyncrl_remote.yaml b/tests/tests_utils/mini_experiment_conf/asyncrl_remote.yaml new file mode 100644 index 0000000..ec1c0bb --- /dev/null +++ b/tests/tests_utils/mini_experiment_conf/asyncrl_remote.yaml @@ -0,0 +1,66 @@ +imports: + - tests.tests_utils.mini_experiment + +engine: asyncrl +debug: false +verbose: false + +env: + name: mini_counter_async + backend: mini_counter + task: MiniCounter-v0 + num_workers: 1 + num_envs: 1 + max_episode_steps: 4 + +buffer: + type: ReplayBuffer + capacity: 128 + sampler: + type: uniform + storage: + type: unified + unit: transition + device: cpu + +policy: + type: MiniAsyncPolicy + rollout_mode: async + train_num_gpus: 0.0 + eval_num_gpus: 0.0 + weight_buffer: + type: WeightBuffer + buffer_strategy: Double + optim_cfg: + lr: 0.05 + +train: + max_epochs: 1 + max_rollout_steps: 4 + batch_size: 8 + eval_interval: -1 + save_interval: -1 + inner_updates: 2 + +cluster: + ray_address: local + train_worker_num: 1 + eval_worker_num: 1 + train_each_gpu_num: 0.0 + eval_each_gpu_num: 0.0 + buffer_worker_num: 1 + remote_train: true + remote_eval: true + remote_storage: false + remote_env: false + placement: + mode: auto + strategy: default + env_strategy: default + +log: + backend: tensorboard + level: INFO + log_dir: ./runs + project: tests + name: mini_asyncrl_remote diff --git a/tests/tests_utils/mini_experiment_conf/syncrl_local.yaml b/tests/tests_utils/mini_experiment_conf/syncrl_local.yaml new file mode 100644 index 0000000..01fe00e --- /dev/null +++ b/tests/tests_utils/mini_experiment_conf/syncrl_local.yaml @@ -0,0 +1,50 @@ +imports: + - tests.tests_utils.mini_experiment + +engine: syncrl +debug: false +verbose: false + +env: + name: mini_counter_local + backend: mini_counter + task: MiniCounter-v0 + num_workers: 1 + num_envs: 1 + max_episode_steps: 4 + +buffer: + type: RolloutBuffer + capacity: 64 + sampler: + type: all + storage: + type: unified + unit: transition + device: cpu + +policy: + type: MiniSyncPolicy + rollout_mode: sync + train_num_gpus: 0.0 + eval_num_gpus: 0.0 + weight_buffer: + type: WeightBuffer + buffer_strategy: Double + optim_cfg: + lr: 0.05 + +train: + max_epochs: 1 + max_rollout_steps: 4 + batch_size: 4 + eval_interval: -1 + save_interval: -1 + inner_updates: 2 + +log: + backend: tensorboard + level: INFO + log_dir: ./runs + project: tests + name: mini_syncrl_local diff --git a/tests/tests_utils/mini_experiment_conf/syncrl_remote.yaml b/tests/tests_utils/mini_experiment_conf/syncrl_remote.yaml new file mode 100644 index 0000000..96f477f --- /dev/null +++ b/tests/tests_utils/mini_experiment_conf/syncrl_remote.yaml @@ -0,0 +1,66 @@ +imports: + - tests.tests_utils.mini_experiment + +engine: syncrl +debug: false +verbose: false + +env: + name: mini_counter_remote + backend: mini_counter + task: MiniCounter-v0 + num_workers: 1 + num_envs: 1 + max_episode_steps: 4 + +buffer: + type: RolloutBuffer + capacity: 64 + sampler: + type: all + storage: + type: unified + unit: transition + device: cpu + +policy: + type: MiniSyncPolicy + rollout_mode: sync + train_num_gpus: 0.0 + eval_num_gpus: 0.0 + weight_buffer: + type: WeightBuffer + buffer_strategy: Double + optim_cfg: + lr: 0.05 + +train: + max_epochs: 1 + max_rollout_steps: 4 + batch_size: 4 + eval_interval: -1 + save_interval: -1 + inner_updates: 2 + +cluster: + ray_address: local + train_worker_num: 1 + eval_worker_num: 1 + train_each_gpu_num: 0.0 + eval_each_gpu_num: 0.0 + buffer_worker_num: 1 + remote_train: true + remote_eval: true + remote_storage: false + remote_env: false + placement: + mode: auto + strategy: default + env_strategy: default + +log: + backend: tensorboard + level: INFO + log_dir: ./runs + project: tests + name: mini_syncrl_remote diff --git a/tests/tests_utils/remote_class_helper.py b/tests/tests_utils/remote_class_helper.py new file mode 100644 index 0000000..3af7bc4 --- /dev/null +++ b/tests/tests_utils/remote_class_helper.py @@ -0,0 +1,6 @@ +from rlightning.utils.ray.remote_class import RayActorMixin + + +class Custom(RayActorMixin): + def __init__(self): + super().__init__() diff --git a/tests/tests_utils/run_mini_experiment.py b/tests/tests_utils/run_mini_experiment.py new file mode 100644 index 0000000..f2ad654 --- /dev/null +++ b/tests/tests_utils/run_mini_experiment.py @@ -0,0 +1,41 @@ +"""Launch test-only minimal experiments through the normal RLightning entrypoint.""" + +from __future__ import annotations + +from pathlib import Path + +import torch + +from rlightning.utils.builders import build_data_buffer, build_engine, build_env_group, build_policy_group +from rlightning.utils.config import MainConfig +from rlightning.utils.launch import launch + + +def main(config: MainConfig) -> None: + env_group = build_env_group(config.env) + policy_group = build_policy_group( + policy_cls=config.policy.type, + policy_cfg=config.policy, + cluster_cfg=config.cluster, + backend="nccl" if torch.cuda.is_available() else "gloo", + ) + buffer = build_data_buffer( + buffer_cls=config.buffer.type, + buffer_cfg=config.buffer, + ) + engine = build_engine( + config=config, + env_group=env_group, + policy_group=policy_group, + buffer=buffer, + ) + + try: + engine.run() + finally: + policy_group.shutdown() + env_group.close() + + +if __name__ == "__main__": + launch(main_func=main, config_path=Path(__file__).parent / "mini_experiment_conf") diff --git a/tests/unit/test_async_rl_engine.py b/tests/unit/test_async_rl_engine.py new file mode 100644 index 0000000..ec20db9 --- /dev/null +++ b/tests/unit/test_async_rl_engine.py @@ -0,0 +1,265 @@ +""" +Unit tests for AsyncRLEngine. + +Tests the core functionality of the async RL engine including initialization, +warm-up, and basic engine operations without external dependencies. +""" + +from unittest.mock import MagicMock, patch + +import pytest + +from rlightning.engine.async_rl_engine import AsyncRLEngine + + + +class _SimpleMocker: + def __init__(self) -> None: + self._patchers = [] + + def Mock(self, *args, **kwargs): + return MagicMock(*args, **kwargs) + + def patch(self, target, *args, **kwargs): + patcher = patch(target, *args, **kwargs) + mocked = patcher.start() + self._patchers.append(patcher) + return mocked + + def stopall(self) -> None: + while self._patchers: + self._patchers.pop().stop() + + +@pytest.fixture +def mocker(): + helper = _SimpleMocker() + try: + yield helper + finally: + helper.stopall() + + +@pytest.fixture(autouse=True) +def debug_mode_env(monkeypatch): + monkeypatch.setenv("RLIGHTNING_DEBUG", "1") + + +@pytest.fixture +def mock_config(mocker): + config = mocker.Mock() + config.train.max_epochs = 2 + config.train.max_rollout_steps = 10 + config.train.batch_size = 4 + return config + + +@pytest.fixture +def mock_env_group(mocker): + env_group = mocker.Mock() + env_group.observation_space.return_value = mocker.Mock() + env_group.action_space.return_value = mocker.Mock() + # warm_up uses reset() and step() returning (requests, truncations) + env_group.reset.return_value = (mocker.Mock(), []) + env_group.step.return_value = (mocker.Mock(), []) + env_group.init.return_value = (mocker.Mock(), {}) + return env_group + + +@pytest.fixture +def mock_policy_group(mocker): + policy_group = mocker.Mock() + policy_group.init_eval.return_value = mocker.Mock() + policy_group.train_list = [mocker.Mock()] + return policy_group + + +@pytest.fixture +def mock_buffer(mocker): + buffer = mocker.Mock() + buffer.size.return_value = 10 # Always has enough data + return buffer + + +def test_engine_initialization(mock_config, mock_env_group, mock_policy_group, mock_buffer): + engine = AsyncRLEngine( + config=mock_config, + env_group=mock_env_group, + policy_group=mock_policy_group, + buffer=mock_buffer, + ) + + assert engine.config == mock_config + assert engine.env_group == mock_env_group + assert engine.policy_group == mock_policy_group + assert engine.buffer == mock_buffer + + +def test_warm_up_calls_initialization_methods(mock_config, mock_env_group, mock_policy_group, mock_buffer): + AsyncRLEngine( + config=mock_config, + env_group=mock_env_group, + policy_group=mock_policy_group, + buffer=mock_buffer, + ) + + mock_env_group.init.assert_called_once() + mock_policy_group.init_eval.assert_called_once() + mock_policy_group.init_train.assert_called_once() + mock_buffer.init.assert_called_once() + + +def test_warm_up_performs_dummy_rollout(mock_config, mock_env_group, mock_policy_group, mock_buffer, mocker): + mock_requests = mocker.Mock() + mock_responses = mocker.Mock() + mock_responses.ids.return_value = ["episode_1"] + + mock_env_group.reset.return_value = (mock_requests, []) + mock_env_group.step.return_value = (mock_requests, []) + mock_policy_group.rollout_batch.return_value = mock_responses + + AsyncRLEngine( + config=mock_config, + env_group=mock_env_group, + policy_group=mock_policy_group, + buffer=mock_buffer, + ) + + assert mock_env_group.reset.call_count >= 1 + assert mock_policy_group.rollout_batch.call_count >= 10 + assert mock_buffer.add_batched_data_async.call_count >= 20 + + +def test_warm_up_performs_dummy_training(mock_config, mock_env_group, mock_policy_group, mock_buffer): + AsyncRLEngine( + config=mock_config, + env_group=mock_env_group, + policy_group=mock_policy_group, + buffer=mock_buffer, + ) + + mock_policy_group.update_dataset.assert_called_once() + mock_policy_group.train.assert_called_once() + mock_policy_group.sync_weights.assert_called_once() + + +def test_engine_initialization_without_debug_warmup( + mock_config, mock_env_group, mock_policy_group, mock_buffer, mocker, monkeypatch +): + monkeypatch.setenv("RLIGHTNING_DEBUG", "0") + mock_policy_group.send_weights = mocker.Mock() + mock_policy_group.notify_update_weights = mocker.Mock() + + warm_up_mock = mocker.patch("rlightning.engine.async_rl_engine.AsyncRLEngine.warm_up") + + AsyncRLEngine( + config=mock_config, + env_group=mock_env_group, + policy_group=mock_policy_group, + buffer=mock_buffer, + ) + + warm_up_mock.assert_not_called() + mock_env_group.init.assert_called_once() + mock_policy_group.init_eval.assert_called_once() + mock_policy_group.init_train.assert_called_once() + mock_buffer.init.assert_called_once() + mock_policy_group.sync_weights.assert_called_once() + + +def test_rollout_loop_structure(mock_config, mock_env_group, mock_policy_group, mock_buffer): + mock_requests = MagicMock() + mock_responses = MagicMock() + + mock_env_group.reset.return_value = (mock_requests, []) + mock_env_group.step_async.return_value = None + mock_env_group.collect_async.return_value = (mock_requests, []) + mock_policy_group.rollout_batch.return_value = mock_responses + + engine = AsyncRLEngine( + config=mock_config, + env_group=mock_env_group, + policy_group=mock_policy_group, + buffer=mock_buffer, + ) + + assert hasattr(engine, "rollout") + assert callable(engine.rollout) + + +def test_train_loop_structure(mock_config, mock_env_group, mock_policy_group, mock_buffer): + engine = AsyncRLEngine( + config=mock_config, + env_group=mock_env_group, + policy_group=mock_policy_group, + buffer=mock_buffer, + ) + + assert hasattr(engine, "train") + assert callable(engine.train) + + +def test_update_weights_loop_structure(mock_config, mock_env_group, mock_policy_group, mock_buffer): + engine = AsyncRLEngine( + config=mock_config, + env_group=mock_env_group, + policy_group=mock_policy_group, + buffer=mock_buffer, + ) + + assert hasattr(engine, "sync_weights") + assert callable(engine.sync_weights) + + +def test_run_method_starts_threads(mock_config, mock_env_group, mock_policy_group, mock_buffer, mocker): + mock_thread = mocker.patch("threading.Thread") + mock_thread_instance = mocker.Mock() + mock_thread.return_value = mock_thread_instance + + engine = AsyncRLEngine( + config=mock_config, + env_group=mock_env_group, + policy_group=mock_policy_group, + buffer=mock_buffer, + ) + + engine.coordinator._done_event.set() + + engine.run() + assert mock_thread.call_count == 4 + assert mock_thread_instance.start.call_count == 4 + assert mock_thread_instance.join.call_count == 4 + + +@pytest.mark.parametrize( + "max_epochs,max_rollout_steps,batch_size", + [ + (1, 5, 2), + (3, 15, 8), + (10, 100, 32), + ], +) +def test_engine_with_different_configs( + max_epochs, + max_rollout_steps, + batch_size, + mock_env_group, + mock_policy_group, + mock_buffer, +): + config = MagicMock() + config.train.max_epochs = max_epochs + config.train.max_rollout_steps = max_rollout_steps + config.train.batch_size = batch_size + + engine = AsyncRLEngine( + config=config, + env_group=mock_env_group, + policy_group=mock_policy_group, + buffer=mock_buffer, + ) + + assert engine.config.train.max_epochs == max_epochs + assert engine.config.train.max_rollout_steps == max_rollout_steps + assert engine.config.train.batch_size == batch_size + assert engine.config.train.batch_size == batch_size diff --git a/tests/unit/test_base_buffer.py b/tests/unit/test_base_buffer.py new file mode 100644 index 0000000..8c9a30a --- /dev/null +++ b/tests/unit/test_base_buffer.py @@ -0,0 +1,776 @@ +""" +Planned unit tests for DataBuffer in base_buffer.py. +""" + +from unittest import mock + +import numpy as np +import pytest + +from rlightning.buffer.base_buffer import DataBuffer +from rlightning.buffer.sampler import AllDataSampler, BatchSampler, UniformSampler +from rlightning.buffer.utils.storage import Storage +from rlightning.buffer.utils.table import EpisodeTable +from rlightning.types.batched_data import BatchedData +from rlightning.utils.config import BufferConfig + +def _make_buffer_config( + sampler_type: str = "uniform", storage_type: str = "unified", num_shards: int = 2 +) -> BufferConfig: + storage_cfg = { + "type": storage_type, + "device": "cpu", + "unit": "transition", + } + if storage_type == "sharded": + storage_cfg["num_shards"] = num_shards + return BufferConfig.from_dict( + { + "type": "ReplayBuffer", + "capacity": 10, + "sampler": {"type": sampler_type}, + "storage": storage_cfg, + } + ) + + +class _FakeSubmitter: + def __init__(self): + self.calls = [] + + def submit(self, method, *args, _block=False, **kwargs): + self.calls.append((method.__name__, args, kwargs, _block)) + return method(*args, **kwargs) + + +class _FakeStorage: + def __init__(self, size: int = 0, label: str = "storage"): + self.size = size + self.label = label + self.data = [] + self.cleared = False + self.env_stats = {} + + def add_episode(self, episode, num_envs): + self.data.append(("add_episode", episode, num_envs)) + + def truncate_one_episode(self, item): + self.data.append(("truncate_one_episode", item)) + + def truncate_episodes(self, items): + self.data.append(("truncate_episodes", list(items))) + + def add_transition(self, env_ret, policy_resp): + self.data.append(("add_transition", env_ret, policy_resp)) + + def add_data_async(self, data): + self.data.append(("add_data_async", data)) + + def get_size(self): + return self.size + + def __getitem__(self, item): + return (self.label, item) + + def get_data(self): + return self.data + + def get_env_stats(self, reset): + return self.env_stats + + def clear(self): + self.cleared = True + self.size = 0 + + def print_timing_summary(self, reset): + self.data.append(("print_timing_summary", reset)) + + +class _FakeSampler: + def __init__(self, indices): + self.indices = indices + + def sample(self, batch_size, data_size, shuffle=True): + if batch_size is None: + return list(range(data_size)) + return list(self.indices)[:batch_size] + + +def test_sanity_check_preprocessor_matrix(): + """Ensure _sanity_check follows the preprocess conflict matrix. + There are three levels of preprocessors, and _sanity_check should + ensure the configuration and preprocess functions are legal. + """ + + cfg = _make_buffer_config() + + def _custom_preprocessor(_): + return _ + + def _custom_env_ret_preprocess_fn(_env_ret): + return _env_ret + + def _custom_preprocess_fn(_env_ret, _policy_resp): + return _env_ret, _policy_resp + + cases = [ + # Default (Happy Path) + ({}, False), + # Only atomic-level custom (Obs/Reward) + ({"obs_preprocessor": _custom_preprocessor}, False), + # Only component-level custom (EnvRet) + ({"env_ret_preprocess_fn": _custom_env_ret_preprocess_fn}, False), + # Only global-level custom (Global preprocess_fn) + ({"preprocess_fn": _custom_preprocess_fn}, False), + # Conflict: atomic vs component + ( + { + "obs_preprocessor": _custom_preprocessor, + "env_ret_preprocess_fn": _custom_env_ret_preprocess_fn, + }, + True, + ), + # Conflict: atomic vs global + ( + { + "obs_preprocessor": _custom_preprocessor, + "preprocess_fn": _custom_preprocess_fn, + }, + True, + ), + # Conflict: component vs global + ( + { + "env_ret_preprocess_fn": _custom_env_ret_preprocess_fn, + "preprocess_fn": _custom_preprocess_fn, + }, + True, + ), + ] + + for kwargs, should_raise in cases: + if should_raise: + with pytest.raises(ValueError): + _ = DataBuffer(config=cfg, **kwargs) + else: + _ = DataBuffer(config=cfg, **kwargs) + + +def test_init_sampler(): + """Plan: build configs for "uniform", "all", "batch" samplers and + verify DataBuffer._init_sampler instantiates the expected sampler class. + Unknown sampler types should raise ValueError, but it will be privented + by earlier config validationin from_dict when we try to make buffer config, + so we skip that case here. + """ + + cfg_uniform = _make_buffer_config(sampler_type="uniform") + buf_uniform = DataBuffer(config=cfg_uniform) + buf_uniform._init_sampler() + assert isinstance(buf_uniform.sampler, UniformSampler) + + cfg_all = _make_buffer_config(sampler_type="all") + buf_all = DataBuffer(config=cfg_all) + buf_all._init_sampler() + assert isinstance(buf_all.sampler, AllDataSampler) + + cfg_batch = _make_buffer_config(sampler_type="batch") + buf_batch = DataBuffer(config=cfg_batch) + buf_batch._init_sampler() + assert isinstance(buf_batch.sampler, BatchSampler) + + +# In base_buffer.py, _init_storage determines where to store data (locally or remotely) and how +# to organize it (unified or sharded). Need to be tested. +# Once local built, storages should be Storage instances, and table should be EpisodeTable with num_storages=1. +def test_init_storage_unified_local_builds_storage(monkeypatch): + """locally storage, verify storages is one Storage instance, table is None.""" + # Plan: force InternalFlag.REMOTE_STORAGE=False, call init(), + # then assert storages[0] is a Storage instance and table is EpisodeTable with num_storages=1. + monkeypatch.setenv("RLIGHTNING_REMOTE_STORAGE", "0") + cfg = _make_buffer_config(storage_type="unified") + buffer = DataBuffer(config=cfg) + buffer.init(env_meta_list=None, env_ids=None) + + assert list(buffer.storages.keys()) == [0] + assert isinstance(buffer.storages[0], Storage) + assert buffer.table.num_storages == 1 + + +# Once remote built, unified storage should have one storage actor and EpisodeTable. +# And the table shouldn't be None. +def test_init_storage_sharded_remote_builds_episode_table(monkeypatch): + """remote sharded storage, verify storages has num_shards entries, table is EpisodeTable.""" + # Plan: force InternalFlag.REMOTE_STORAGE=True, stub ray.remote to a fake, + # then assert storages has num_shards entries and table is EpisodeTable. + monkeypatch.setenv("RLIGHTNING_REMOTE_STORAGE", "1") + cfg = _make_buffer_config(storage_type="sharded", num_shards=2) + buffer = DataBuffer(config=cfg) + + _mock_grm = mock.MagicMock() + _mock_scheduling = mock.MagicMock() + _mock_scheduling.buffer_worker.worker_num = 2 + _mock_scheduling.train_worker.worker_num = 2 + _mock_grm.get_scheduling.return_value = _mock_scheduling + + buffer._global_resource_manager = _mock_grm + + def _fake_init_storage_actor(self, env_meta_list, gpu_requirement, scheduling_kwargs): + return f"storage-{len(self.storages)}" + + buffer._init_storage_actor = _fake_init_storage_actor.__get__(buffer, DataBuffer) + buffer.init(env_meta_list=None, env_ids=["env-1", "env-2"]) + + assert set(buffer.storages.keys()) == {0, 1} + assert isinstance(buffer.table, EpisodeTable) + assert buffer.table.num_storages == 2 + + +# in fact, still need to cover unified+remote and sharded+local and unknown type. if I had time ... + +# test add_transition + + +def test_add_transition_truncate_calls_storage_methods(): + """truncated=true""" + # Plan: inject fake storage + TaskSubmitter, call add_transition with truncated=True, + # then assert add_transition and truncate_one_episode are called in order. + cfg = _make_buffer_config() + buffer = DataBuffer(config=cfg) + buffer.storages = {0: _FakeStorage()} + buffer.task_submitter = _FakeSubmitter() + buffer.init_storage_table() + + buffer.add_transition("env-1", {"r": 1}, {"a": 0}, truncated=True) + + assert [call[0] for call in buffer.task_submitter.calls] == [ + "add_transition", + "truncate_one_episode", + ] + assert buffer.task_submitter.calls[0][-1] is False + assert buffer.task_submitter.calls[1][-1] is False + + +def test_add_transition_routes_to_sharded_storage(): + """sharded storage should route by env_id.""" + # Plan: set EpisodeTable with two envs, ensure env-2 routes to storage 1. + cfg = _make_buffer_config() + buffer = DataBuffer(config=cfg) + buffer.storages = {0: _FakeStorage(label="s0"), 1: _FakeStorage(label="s1")} + buffer.init_storage_table(env_ids=["env-1", "env-2"], train_worker_num=2) + buffer.task_submitter = _FakeSubmitter() + + buffer.add_transition("env-2", {"r": 1}, {"a": 0}, truncated=False) + + assert buffer.storages[1].data[0][0] == "add_transition" + assert buffer.storages[0].data == [] + + +def test_add_batched_transition_validates_length_and_ids(): + """whether the original function could raise ValueErorr or TypeError when lengths or + ids mismatch. + """ + # Plan: create fake BatchedData with mismatched length/ids and assert ValueError. + cfg = _make_buffer_config() + buffer = DataBuffer(config=cfg) + + env_ret = BatchedData(ids=["env-1"], data=[{"r": 1}]) + policy_resp = BatchedData(ids=["env-1", "env-2"], data=[{"a": 0}, {"a": 1}]) + with pytest.raises(ValueError): + buffer.add_batched_transition(env_ret, policy_resp) + + env_ret = BatchedData(ids=["env-1", "env-2"], data=[{"r": 1}, {"r": 2}]) + policy_resp = BatchedData(ids=["env-1", "env-3"], data=[{"a": 0}, {"a": 1}]) + with pytest.raises(ValueError): + buffer.add_batched_transition(env_ret, policy_resp) + + +def test_add_batched_transition_validates_truncations(): + """whether the original function could raise ValueErorr or TypeError when truncations + is invalid. + """ + cfg = _make_buffer_config() + buffer = DataBuffer(config=cfg) + + env_ret = BatchedData(ids=["env-1", "env-2"], data=[{"r": 1}, {"r": 2}]) + policy_resp = BatchedData(ids=["env-1", "env-2"], data=[{"a": 0}, {"a": 1}]) + + with pytest.raises(TypeError): + buffer.add_batched_transition(env_ret, policy_resp, truncations="not-a-list") + + with pytest.raises(ValueError): + buffer.add_batched_transition(env_ret, policy_resp, truncations=[True]) + + +def test_add_batched_transition_calls_add_transition_and_returns_self(): + """default truncations=None should call add_transition for each item.""" + cfg = _make_buffer_config() + buffer = DataBuffer(config=cfg) + buffer.storages = {0: _FakeStorage()} + buffer.init_storage_table() + buffer.task_submitter = _FakeSubmitter() + + env_ret = BatchedData(ids=["env-1", "env-2"], data=[{"r": 1}, {"r": 2}]) + policy_resp = BatchedData(ids=["env-1", "env-2"], data=[{"a": 0}, {"a": 1}]) + + result = buffer.add_batched_transition(env_ret, policy_resp, truncations=None) + + assert result is buffer + assert [call[0] for call in buffer.task_submitter.calls] == [ + "add_transition", + "add_transition", + ] + assert buffer.task_submitter.calls[0][1] == ({"r": 1}, {"a": 0}) + assert buffer.task_submitter.calls[1][1] == ({"r": 2}, {"a": 1}) + + +def test_add_batched_transition_passes_truncations_per_item(): + """truncations per item should pass to add_transition calls.""" + cfg = _make_buffer_config() + buffer = DataBuffer(config=cfg) + buffer.storages = {0: _FakeStorage()} + buffer.init_storage_table() + buffer.task_submitter = _FakeSubmitter() + + env_ret = BatchedData(ids=["env-1", "env-2"], data=[{"r": 1}, {"r": 2}]) + policy_resp = BatchedData(ids=["env-1", "env-2"], data=[{"a": 0}, {"a": 1}]) + + buffer.add_batched_transition(env_ret, policy_resp, truncations=[True, False]) + + transition_calls = [call for call in buffer.task_submitter.calls if call[0] == "add_transition"] + assert transition_calls[0][1] == ({"r": 1}, {"a": 0}) + assert transition_calls[1][1] == ({"r": 2}, {"a": 1}) + assert buffer.task_submitter.calls[0][-1] is False + assert buffer.task_submitter.calls[1][-1] is False + assert buffer.storages[0].data[1][0] == "truncate_one_episode" + + +def test_truncate_one_episode_rejects_invalid_item(): + """test if env_id is None, raise TypeError.""" + # Plan: call truncate_one_episode with object missing env_id and assert TypeError. + cfg = _make_buffer_config() + buffer = DataBuffer(config=cfg) + + class _NoEnvId: + pass + + with pytest.raises(TypeError): + buffer.truncate_one_episode(_NoEnvId()) + + +def test_truncate_one_episode_routes_and_passes_item(): + """test routes by using different env_ids. and test passing object with env_id attr.""" + cfg = _make_buffer_config(storage_type="sharded", num_shards=2) + buffer = DataBuffer(config=cfg) + buffer.storages = {0: _FakeStorage(label="s0"), 1: _FakeStorage(label="s1")} + buffer.init_storage_table(env_ids=["env-1", "env-2"], train_worker_num=2) + buffer.task_submitter = _FakeSubmitter() + + # 1. direct env_id string + buffer.truncate_one_episode("env-1") + + class _Item: + env_id = "env-2" + + obj = _Item() + # 2. object with env_id attribute + buffer.truncate_one_episode(obj) + + assert ("truncate_one_episode", "env-1") in buffer.storages[0].data + assert ("truncate_one_episode", obj) in buffer.storages[1].data + + +def test_get_rejects_invalid_item_type(): + """object() will raise TypeError in get().""" + # Plan: call get() with a non-dict, non-int, non-sequence and assert TypeError. + cfg = _make_buffer_config() + buffer = DataBuffer(config=cfg) + + with pytest.raises(TypeError): + buffer.get(object()) + + # input None + with pytest.raises(TypeError): + buffer.get(None) + + # input float + with pytest.raises(TypeError): + buffer.get(3.14) + + # TODO: input dict but wrong items + + +def test_len_and_size_sum_storage_sizes(): + """verify len() and size() sum across storages and query each shard.""" + cfg = _make_buffer_config() + buffer = DataBuffer(config=cfg) + buffer.storages = {0: _FakeStorage(size=3), 1: _FakeStorage(size=5)} + buffer.task_submitter = _FakeSubmitter() + + # and __len__ in base_buffer.py use _block=True and submit to task_submitter, + # so... + assert len(buffer) == 8 + assert len(buffer.task_submitter.calls) == 2 + assert all(call[0] == "get_size" for call in buffer.task_submitter.calls) + assert all(call[-1] is True for call in buffer.task_submitter.calls) + + buffer.task_submitter = _FakeSubmitter() + assert buffer.size() == 8 + assert len(buffer.task_submitter.calls) == 2 + assert all(call[0] == "get_size" for call in buffer.task_submitter.calls) + assert all(call[-1] is True for call in buffer.task_submitter.calls) + + +def test_add_episode_uses_random_storage(monkeypatch): + """test add_episode uses random.choice to select storage. + but here we fix the choice to always return storage 1. + """ + cfg = _make_buffer_config() + buffer = DataBuffer(config=cfg) + buffer.storages = {0: _FakeStorage(label="s0"), 1: _FakeStorage(label="s1")} + buffer.task_submitter = _FakeSubmitter() + + monkeypatch.setattr("rlightning.buffer.base_buffer.random.choice", lambda keys: 1) + buffer.add_episode({"ep": 1}, num_envs=1) + + assert buffer.storages[1].data[0][0] == "add_episode" + # so far we still need to verify what we had storaged. + call_record = buffer.storages[1].data[0] + assert call_record[0] == "add_episode" + assert call_record[1] == {"ep": 1} + assert call_record[2] == 1 + + +# the point of add_data_async is the "async". but how to test it? +# truncated is true or false, different env_ids, and keep data unchanged. + + +def test_add_data_async_truncate_calls_storage_methods(): + """Validate behavior of add_data_async when truncated=True. + Specifically, ensure: + 1. tasks are submitted in the correct order, and + 2. submissions are non-blocking (i.e. _block=False). + """ + cfg = _make_buffer_config() + buffer = DataBuffer(config=cfg) + buffer.storages = {0: _FakeStorage()} + buffer.init_storage_table() + buffer.task_submitter = _FakeSubmitter() + + data = {"r": 1} + buffer.add_data_async("env-1", data, truncated=True) + + assert [call[0] for call in buffer.task_submitter.calls] == [ + "add_data_async", + "truncate_one_episode", + ] + assert buffer.task_submitter.calls[0][1][0] == data + assert buffer.task_submitter.calls[1][1][0] == data + assert buffer.task_submitter.calls[0][-1] is False + assert buffer.task_submitter.calls[1][-1] is False + + +def test_add_data_async_no_truncate_only_adds(): + """Validate behavior of add_data_async when truncated=False. + Specifically, ensure only add_data_async is called. + """ + cfg = _make_buffer_config() + buffer = DataBuffer(config=cfg) + buffer.storages = {0: _FakeStorage()} + buffer.init_storage_table() + buffer.task_submitter = _FakeSubmitter() + + data = {"r": 2} + buffer.add_data_async("env-1", data, truncated=False) + + assert [call[0] for call in buffer.task_submitter.calls] == ["add_data_async"] + assert buffer.task_submitter.calls[0][1][0] == data + assert buffer.task_submitter.calls[0][-1] is False + + +def test_add_data_async_routes_to_sharded_storage(): + """sharded storage should route by env_id. + very import test.""" + cfg = _make_buffer_config() + buffer = DataBuffer(config=cfg) + buffer.storages = {0: _FakeStorage(label="s0"), 1: _FakeStorage(label="s1")} + buffer.init_storage_table(env_ids=["env-1", "env-2"], train_worker_num=2) + buffer.task_submitter = _FakeSubmitter() + + data = {"r": 3} + buffer.add_data_async("env-2", data, truncated=False) + + assert buffer.storages[1].data[0] == ("add_data_async", data) + assert buffer.storages[0].data == [] + + +def test_add_batched_data_async_validates_truncations(): + """as you can see, validation of truncations in add_batched_data_async.""" + cfg = _make_buffer_config() + buffer = DataBuffer(config=cfg) + data = BatchedData(ids=["env-1", "env-2"], data=[{"r": 1}, {"r": 2}]) + + with pytest.raises(TypeError): + buffer.add_batched_data_async(data, truncations="not-a-list") + + with pytest.raises(ValueError): + buffer.add_batched_data_async(data, truncations=[True]) + + +def test_add_batched_data_async_default_truncations_calls_add_data_async(): + """test default truncations=None calls add_data_async for each item.""" + cfg = _make_buffer_config() + buffer = DataBuffer(config=cfg) + buffer.storages = {0: _FakeStorage()} + buffer.init_storage_table() + buffer.task_submitter = _FakeSubmitter() + data = BatchedData(ids=["env-1", "env-2"], data=[{"r": 1}, {"r": 2}]) + + buffer.add_batched_data_async(data, truncations=None) + + assert [call[0] for call in buffer.task_submitter.calls] == [ + "add_data_async", + "add_data_async", + ] + assert buffer.task_submitter.calls[0][1][0] == {"r": 1} + assert buffer.task_submitter.calls[1][1][0] == {"r": 2} + assert buffer.task_submitter.calls[0][-1] is False + assert buffer.task_submitter.calls[1][-1] is False + + +def test_add_batched_data_async_passes_truncations_per_item(): + """batch add two data, one truncated, one not. + test if truncations per item passed to add_data_async calls.""" + cfg = _make_buffer_config() + buffer = DataBuffer(config=cfg) + buffer.storages = {0: _FakeStorage()} + buffer.init_storage_table() + buffer.task_submitter = _FakeSubmitter() + data = BatchedData(ids=["env-1", "env-2"], data=[{"r": 1}, {"r": 2}]) + + buffer.add_batched_data_async(data, truncations=[True, False]) + + add_calls = [call for call in buffer.task_submitter.calls if call[0] == "add_data_async"] + assert len(add_calls) == 2 + assert add_calls[0][1][0] == {"r": 1} + assert add_calls[1][1][0] == {"r": 2} + assert buffer.storages[0].data[1][0] == "truncate_one_episode" + + +def test_truncate_episodes_groups_by_storage(): + """test truncate_episodes groups by storage using env_ids. + group route by storage index.""" + cfg = _make_buffer_config(storage_type="sharded", num_shards=2) + buffer = DataBuffer(config=cfg) + buffer.table = EpisodeTable(num_storages=2) + buffer.storages = {0: _FakeStorage(label="s0"), 1: _FakeStorage(label="s1")} + buffer.init_storage_table(env_ids=["env-1", "env-2"], train_worker_num=2) + buffer.task_submitter = _FakeSubmitter() + + buffer.truncate_episodes(["env-1", "env-2"]) + + assert ("truncate_episodes", ["env-1"]) in buffer.storages[0].data + assert ("truncate_episodes", ["env-2"]) in buffer.storages[1].data + + +def test_truncate_episodes_accepts_objects_and_routes_by_env_id(): + """test truncate_episodes accepts objects and routes by env_id.""" + cfg = _make_buffer_config(storage_type="sharded", num_shards=2) + buffer = DataBuffer(config=cfg) + buffer.table = EpisodeTable(num_storages=2, env_ids=["env-1", "env-2"]) + buffer.storages = {0: _FakeStorage(label="s0"), 1: _FakeStorage(label="s1")} + buffer.init_storage_table(env_ids=["env-1", "env-2"], train_worker_num=2) + buffer.task_submitter = _FakeSubmitter() + + class _Item: + def __init__(self, env_id): + self.env_id = env_id + + buffer.truncate_episodes([_Item("env-1"), _Item("env-2")]) + + assert ("truncate_episodes", ["env-1"]) in buffer.storages[0].data + assert ("truncate_episodes", ["env-2"]) in buffer.storages[1].data + + +def test_truncate_episodes_rejects_invalid_item(): + """invalid test""" + cfg = _make_buffer_config(storage_type="sharded", num_shards=2) + buffer = DataBuffer(config=cfg) + buffer.table = EpisodeTable(num_storages=2) + buffer.storages = {0: _FakeStorage(label="s0"), 1: _FakeStorage(label="s1")} + buffer.init_storage_table(train_worker_num=2) + buffer.task_submitter = _FakeSubmitter() + + class _NoEnvId: + pass + + with pytest.raises(TypeError): + buffer.truncate_episodes(["env-1", _NoEnvId()]) + + assert buffer.task_submitter.calls == [] + + +def test_getitem_and_get_dict_route_to_storage(): + """test __getitem__ and get method route to storage correctly.""" + cfg = _make_buffer_config() + buffer = DataBuffer(config=cfg) + buffer.storages = {0: _FakeStorage(label="s0")} + buffer.task_submitter = _FakeSubmitter() + + assert buffer[0, [1, 2]] == ("s0", [1, 2]) + assert buffer.get({"storage_idx": 0, "indices": [3, 4]}) == ("s0", [3, 4]) + + +def test_get_all_single_and_multi_storage(): + """test get_all works for single and multi storage setups.""" + cfg = _make_buffer_config() + buffer = DataBuffer(config=cfg) + buffer.storages = {0: _FakeStorage(label="s0")} + buffer.task_submitter = _FakeSubmitter() + + single = buffer.get_all() + assert single == [] + + buffer.storages = {0: _FakeStorage(label="s0"), 1: _FakeStorage(label="s1")} + multi = buffer.get_all() + assert set(multi.keys()) == {0, 1} + + +def test_sample_raises_on_unequal_shard_sizes(): + """test sample raises RuntimeError when shard sizes are unequal""" + cfg = _make_buffer_config(storage_type="sharded", num_shards=2) + buffer = DataBuffer(config=cfg) + buffer.storages = {0: _FakeStorage(size=3), 1: _FakeStorage(size=4)} + buffer.task_submitter = _FakeSubmitter() + buffer.sampler = _FakeSampler(indices=[0, 1, 2]) + buffer.init_storage_table(train_worker_num=2) + buffer.table._storage_to_train_workers == {0: [0], 1: [1]} + + with pytest.raises(RuntimeError): + buffer.sample(batch_size=4) + + +def test_sample_splits_across_workers_drop_last(): + """Covers basic split across shards with drop_last enabled and equal-sized splits.""" + cfg = _make_buffer_config(storage_type="sharded", num_shards=2) + buffer = DataBuffer(config=cfg) + buffer.storages = {0: _FakeStorage(size=4, label="s0"), 1: _FakeStorage(size=4, label="s1")} + buffer.task_submitter = _FakeSubmitter() + buffer.sampler = _FakeSampler(indices=[0, 1, 2, 3]) + buffer.init_storage_table(train_worker_num=2) + buffer.table._storage_to_train_workers == {0: [0], 1: [1]} + + sample_data = buffer.sample(batch_size=4, shuffle=False, drop_last=True) + + assert sample_data[0][0] == "s0" + assert sample_data[0][1].tolist() == [0, 1] + assert sample_data[1][0] == "s1" + assert sample_data[1][1].tolist() == [0, 1] + + +def test_sample_drop_last_truncates_uneven_split(): + """Ensures drop_last truncates indices so each worker gets the same count.""" + cfg = _make_buffer_config(storage_type="unified", num_shards=1) + buffer = DataBuffer(config=cfg) + buffer.storages = {0: _FakeStorage(size=6, label="s0")} + buffer.task_submitter = _FakeSubmitter() + buffer.sampler = _FakeSampler(indices=[0, 1, 2, 3, 4]) + buffer.init_storage_table(train_worker_num=2) + buffer.table._storage_to_train_workers == {0: [0, 1]} + + sample_data = buffer.sample(batch_size=5, shuffle=False, drop_last=True) + + assert sample_data[0][0] == "s0" + assert sample_data[0][1].tolist() == [0, 2] + assert sample_data[1][0] == "s0" + assert sample_data[1][1].tolist() == [1, 3] + + +def test_sample_no_drop_last_pads_uneven_split(): + """Ensures drop_last=False pads indices to make splits even across workers.""" + cfg = _make_buffer_config(storage_type="unified", num_shards=1) + buffer = DataBuffer(config=cfg) + buffer.storages = {0: _FakeStorage(size=6, label="s0")} + buffer.task_submitter = _FakeSubmitter() + buffer.sampler = _FakeSampler(indices=[0, 1, 2, 3, 4]) + buffer.init_storage_table(train_worker_num=2) + buffer.table._storage_to_train_workers == {0: [0, 1]} + + sample_data = buffer.sample(batch_size=5, shuffle=False, drop_last=False) + + assert sample_data[0][0] == "s0" + assert sample_data[0][1].tolist() == [0, 2, 4] + assert sample_data[1][0] == "s0" + assert sample_data[1][1].tolist() == [1, 3, 0] + + +def test_sample_no_drop_last_converts_list_indices_to_array(): + """drop_last=False pads indices; ensure list indices are converted for storage access.""" + cfg = _make_buffer_config(storage_type="unified", num_shards=1) + buffer = DataBuffer(config=cfg) + + class _TypeCheckingStorage(_FakeStorage): + def __getitem__(self, item): + if isinstance(item, list): + raise TypeError("list indices not supported") + return super().__getitem__(item) + + buffer.storages = {0: _TypeCheckingStorage(size=6, label="s0")} + buffer.task_submitter = _FakeSubmitter() + buffer.sampler = _FakeSampler(indices=[0, 1, 2, 3, 4]) + buffer.init_storage_table(train_worker_num=2) + buffer.table._storage_to_train_workers = {0: [0, 1]} + + sample_data = buffer.sample(batch_size=5, shuffle=False, drop_last=False) + + assert isinstance(sample_data[0][1], np.ndarray) + assert isinstance(sample_data[1][1], np.ndarray) + assert sample_data[0][1].tolist() == [0, 2, 4] + assert sample_data[1][1].tolist() == [1, 3, 0] + + +def test_init_storage_table_default_mapping(): + """test init_storage_table default mapping.""" + cfg = _make_buffer_config() + buffer = DataBuffer(config=cfg) + buffer.storages = {0: _FakeStorage()} + + buffer.init_storage_table() + + assert buffer.table._storage_to_train_workers == {0: [0]} + + +def test_init_storage_table_raises_on_manager_error(): + """test init_storage_table raises RuntimeError on manager error.""" + cfg = _make_buffer_config(storage_type="sharded", num_shards=2) + buffer = DataBuffer(config=cfg) + buffer.storages = {0: _FakeStorage(), 1: _FakeStorage()} + + class _BadManager: + def get_component_distribution(self): + raise RuntimeError("bad manager") + + def get_scheduling(self): + raise RuntimeError("bad manager") + + buffer._global_resource_manager = _BadManager() + + with pytest.raises(RuntimeError): + buffer.init_storage_table(train_worker_num=2) + + +def test_clear_calls_storage_clear(): + """test clear calls clear on each storage and resets buffer length.""" + cfg = _make_buffer_config() + buffer = DataBuffer(config=cfg) + s0 = _FakeStorage(size=2) + s1 = _FakeStorage(size=3) + buffer.storages = {0: s0, 1: s1} + buffer.task_submitter = _FakeSubmitter() + + buffer.clear() + + assert s0.cleared is True + assert s1.cleared is True + assert len(buffer) == 0 + assert len(buffer) == 0 diff --git a/tests/unit/test_base_policy.py b/tests/unit/test_base_policy.py new file mode 100644 index 0000000..aa5b6d5 --- /dev/null +++ b/tests/unit/test_base_policy.py @@ -0,0 +1,949 @@ +"""Unit tests for BasePolicy behaviors.""" + +from __future__ import annotations + +import asyncio +import builtins +import threading +from typing import Any, Dict + +import numpy as np +import pytest +import torch +from torch import nn + +from rlightning.policy.base_policy import BasePolicy, PolicyRole +from rlightning.types import EnvRet, PolicyResponse +from rlightning.utils.config import PolicyConfig, TrainConfig +from rlightning.utils.utils import InternalFlag + + +@pytest.fixture(autouse=True) +def _cpu_safe_cuda(monkeypatch): + if torch.cuda.is_available(): + return + + monkeypatch.setattr(EnvRet, "cuda", lambda self, device="cuda": self, raising=False) + monkeypatch.setattr(PolicyResponse, "cuda", lambda self, device="cuda": self, raising=False) + + +class _DummyPolicy(BasePolicy): + """Concrete policy for exercising BasePolicy behaviors.""" + + def __init__(self, config: PolicyConfig, role_type: PolicyRole) -> None: + super().__init__(config, role_type) + self.update_weights_called = False + self.train_called = False + self.loaded_state = None + + def construct_network(self, env_meta: Any = None, *args: Any, **kwargs: Any) -> None: + self.linear = nn.Linear(2, 2) + self.frozen = nn.Linear(2, 2) + for param in self.frozen.parameters(): + param.requires_grad = False + + def setup_optimizer(self, optim_cfg: Any) -> None: + self.optimizer = torch.optim.SGD(self.parameters(), lr=0.1) + + def rollout_step(self, env_ret: EnvRet, **kwargs: Any) -> PolicyResponse: + return PolicyResponse(env_id=env_ret.env_id, action=torch.tensor(1)) + + def postprocess(self, env_ret: EnvRet | None = None, policy_resp: PolicyResponse | None = None) -> Any: + return (env_ret, policy_resp, torch.tensor([1.0])) + + def update_dataset(self, data: Any) -> None: + self.dataset = data + + def train(self, *args: Any, **kwargs: Any) -> Any: + self.train_called = True + return None + + def get_trainable_parameters(self) -> Dict[str, Dict[str, torch.Tensor]]: + params: Dict[str, Dict[str, torch.Tensor]] = {} + for name, model in self.model_list: + module = model.module if hasattr(model, "module") else model + params[name] = module.state_dict() + return params + + def load_state_dict(self, state_dict: Dict[str, torch.Tensor], *args: Any, **kwargs: Any) -> None: + self.loaded_state = state_dict + + def update_weights(self) -> None: + self.update_weights_called = True + return None + + +class _AsyncDummyPolicy(_DummyPolicy): + """Async variant for testing _rollout_async path.""" + + async def rollout_step(self, env_ret: EnvRet, **kwargs: Any) -> PolicyResponse: + return PolicyResponse(env_id=env_ret.env_id, action=torch.tensor(2)) + + +class _WarmupPolicy(_DummyPolicy): + """Policy override to observe warmup rebuild behavior.""" + + def init_train(self, train_config: TrainConfig, env_meta: Any = None) -> None: + self.init_train_called = True + self.is_init = True + + def _setup_sampling_params(self) -> None: + self.setup_sampling_called = True + + +class _FakeLoop: + """Minimal loop stub for call_soon_threadsafe in hook tests.""" + + def __init__(self) -> None: + self.calls = [] + + def call_soon_threadsafe(self, func, *args): + self.calls.append(func) + func(*args) + + +class _FailingRolloutPolicy(_DummyPolicy): + """Policy that raises in rollout_step.""" + + def rollout_step(self, env_ret: EnvRet, **kwargs: Any) -> PolicyResponse: + raise RuntimeError("test rollout error") + + +class _FailingPostprocessPolicy(_DummyPolicy): + """Policy that raises in postprocess.""" + + def postprocess(self, env_ret: EnvRet | None = None, policy_resp: PolicyResponse | None = None) -> Any: + raise RuntimeError("test postprocess error") + + +def _make_policy_config(rollout_mode: str = "sync") -> PolicyConfig: + return PolicyConfig.from_dict( + { + "type": "DummyPolicy", + "rollout_mode": rollout_mode, + } + ) + + +def test_sanity_check_invalid_role_raises(): + """_sanity_check should reject unsupported role_type values.""" + config = _make_policy_config() + with pytest.raises(ValueError): + _ = _DummyPolicy(config, role_type="bad-role") + + +def test_find_model_selects_trainable_modules_only(): + """_find_model should include only modules with requires_grad params.""" + policy = _DummyPolicy(_make_policy_config(), PolicyRole.TRAIN) + policy.construct_network() + policy._find_model() + + names = [name for name, _ in policy.model_list] + assert "linear" in names + assert "frozen" not in names + + +def test_find_model_raises_on_invalid_model_list_entry(): + """_find_model should reject non-nn.Module entries in model_list.""" + policy = _DummyPolicy(_make_policy_config(), PolicyRole.TRAIN) + policy.model_list = [("bad", object())] + + with pytest.raises(ValueError): + policy._find_model() + + +def test_init_eval_sync_sets_idle_event_and_eval_mode(): + """init_eval should set idle event and switch models to eval mode.""" + policy = _DummyPolicy(_make_policy_config("sync"), PolicyRole.EVAL) + policy.init_eval(env_meta=None) + + assert policy.is_init is True + assert policy._idle_as_infer.is_set() is True + assert policy.linear.training is False + + +def test_init_eval_async_sets_async_fields(): + """init_eval should initialize async admission controls and counters.""" + policy = _DummyPolicy(_make_policy_config("async"), PolicyRole.EVAL) + loop = asyncio.new_event_loop() + + try: + try: + old_loop = asyncio.get_event_loop() + except RuntimeError: + old_loop = None + + asyncio.set_event_loop(loop) + policy.init_eval(env_meta=None) + + assert policy.is_init is True + assert policy._loop is loop + assert policy._accept_new_requests.is_set() is True + assert policy.num_requests == 0 + assert isinstance(policy._num_requests_lock, type(threading.Lock())) + assert isinstance(policy._inflight_zero_cv, threading.Condition) + assert policy._update_weights_signal.is_set() is False + finally: + asyncio.set_event_loop(old_loop) + loop.close() + + +def test_init_eval_invalid_rollout_mode_raises(): + """init_eval should reject unknown rollout_mode.""" + policy = _DummyPolicy(_make_policy_config("sync"), PolicyRole.EVAL) + policy.rollout_mode = "invalid" + + with pytest.raises(ValueError): + policy.init_eval(env_meta=None) + + +def test_init_train_sets_train_mode_and_optimizer(): + """init_train should enable train mode and initialize optimizer.""" + policy = _DummyPolicy(_make_policy_config("sync"), PolicyRole.TRAIN) + train_cfg = TrainConfig(max_epochs=1) + + policy.init_train(train_cfg) + + assert policy.is_init is True + assert policy.linear.training is True + assert hasattr(policy, "optimizer") + + +def test_is_initialized_reflects_state(): + """is_initialized should reflect init_train/init_eval state.""" + policy = _DummyPolicy(_make_policy_config("sync"), PolicyRole.TRAIN) + + assert policy.is_initialized() is False + + policy.init_train(TrainConfig(max_epochs=1)) + + assert policy.is_initialized() is True + + +def test_rollout_sync_clears_and_sets_idle_event(): + """_rollout should clear idle flag during execution and set it afterward.""" + + class _StateCapturingPolicy(_DummyPolicy): + """Policy that captures idle state during rollout execution.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.idle_during_rollout = None + + def rollout_step(self, env_ret, **kwargs): + # Capture idle state during rollout execution + self.idle_during_rollout = self._idle_as_infer.is_set() + return super().rollout_step(env_ret, **kwargs) + + policy = _StateCapturingPolicy(_make_policy_config("sync"), PolicyRole.EVAL) + policy._idle_as_infer = threading.Event() + policy._idle_as_infer.set() + + env_ret = EnvRet(env_id="env-1", observation=np.array([1.0])) + resp = policy._rollout(env_ret) + + # Verify idle was cleared during rollout execution + assert policy.idle_during_rollout is False, "idle should be cleared during rollout" + + # Verify idle is set after rollout completes + assert policy._idle_as_infer.is_set() is True, "idle should be set after rollout" + + # Verify response is correct + assert resp.env_id == "env-1" + + +@pytest.mark.parametrize( + "remote_storage,remote_env", + [ + ("1", "0"), # Only REMOTE_STORAGE + ("0", "1"), # Only REMOTE_ENV + ("1", "1"), # Both flags set + ], + ids=["storage_only", "env_only", "both"], +) +def test_post_rollout_hook_converts_when_remote(monkeypatch, remote_storage, remote_env): + """_post_rollout_hook should call numpy on response for any remote flag.""" + policy = _DummyPolicy(_make_policy_config("sync"), PolicyRole.EVAL) + policy._idle_as_infer = threading.Event() + + policy_resp = PolicyResponse(env_id="env-4", action=torch.tensor(1)) + numpy_called = {"called": False} + + def _numpy(): + numpy_called["called"] = True + return policy_resp + + policy_resp.numpy = _numpy + + monkeypatch.setenv("RLIGHTNING_REMOTE_STORAGE", remote_storage) + monkeypatch.setenv("RLIGHTNING_REMOTE_ENV", remote_env) + + result = policy._post_rollout_hook(policy_resp) + + assert ( + numpy_called["called"] is True + ), f"numpy() should be called when REMOTE_STORAGE={remote_storage}, REMOTE_ENV={remote_env}" + assert result is policy_resp + assert policy._idle_as_infer.is_set() is True + + +def test_rollout_async_tracks_inflight_requests(): + """_rollout_async should increment then decrement inflight count.""" + + async def _run(): + class _StateCapturingAsyncPolicy(_AsyncDummyPolicy): + """Policy that captures num_requests during rollout execution.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.num_requests_during_rollout = None + + async def rollout_step(self, env_ret, **kwargs): + # Capture num_requests during rollout execution + self.num_requests_during_rollout = self.num_requests + return await super().rollout_step(env_ret, **kwargs) + + policy = _StateCapturingAsyncPolicy(_make_policy_config("async"), PolicyRole.EVAL) + policy.rollout_mode = "async" + policy._accept_new_requests = asyncio.Event() + policy._accept_new_requests.set() + policy.num_requests = 0 + policy._num_requests_lock = threading.Lock() + policy._inflight_zero_cv = threading.Condition(policy._num_requests_lock) + + env_ret = EnvRet(env_id="env-2", observation=np.array([0.0])) + resp = await policy._rollout_async(env_ret) + return resp, policy.num_requests, policy.num_requests_during_rollout + + resp, num_requests_after, num_requests_during = asyncio.run(_run()) + + # Verify num_requests was incremented during rollout + assert num_requests_during == 1, "num_requests should be 1 during rollout" + + # Verify num_requests is decremented after rollout + assert num_requests_after == 0, "num_requests should be 0 after rollout" + + # Verify response is correct + assert resp.env_id == "env-2" + + +@pytest.mark.parametrize( + "initial_count,should_notify", + [ + (1, True), # 1 → 0, should notify + (2, False), # 2 → 1, should NOT notify + (3, False), # 3 → 2, should NOT notify + ], + ids=["reaches_zero", "still_one", "still_two"], +) +def test_post_rollout_hook_async_notify_behavior(initial_count, should_notify): + """_post_rollout_hook_async should only notify when reaching zero.""" + + class _NotifyCounter: + def __init__(self) -> None: + self.called = False + + def notify_all(self) -> None: + self.called = True + + policy = _DummyPolicy(_make_policy_config("async"), PolicyRole.EVAL) + policy.num_requests = initial_count + policy._num_requests_lock = threading.Lock() + notifier = _NotifyCounter() + policy._inflight_zero_cv = notifier + + policy_resp = PolicyResponse(env_id="env-notify", action=torch.tensor(1)) + result = policy._post_rollout_hook_async(policy_resp) + + assert policy.num_requests == initial_count - 1 + assert notifier.called is should_notify, ( + f"notify_all should{' ' if should_notify else ' NOT '}be called " + f"when num_requests goes from {initial_count} to {initial_count - 1}" + ) + assert result is policy_resp + + +def test_update_weight_hooks_sync_manage_events(): + """_pre/_post_update_weights_hook should manage sync events properly.""" + + class _TrackingEvent: + """Event wrapper that tracks wait() calls.""" + + def __init__(self, name): + self._event = threading.Event() + self._event.set() + self._name = name + self.wait_calls = [] + + def wait(self): + self.wait_calls.append(self._name) + return self._event.wait() + + def set(self): + self._event.set() + + def clear(self): + self._event.clear() + + def is_set(self): + return self._event.is_set() + + policy = _DummyPolicy(_make_policy_config("sync"), PolicyRole.EVAL) + policy.rollout_mode = "sync" + update_signal = _TrackingEvent("update_signal") + idle_event = _TrackingEvent("idle") + policy._update_weights_signal = update_signal + policy._idle_as_infer = idle_event + policy._weight_update_done = None + + policy._pre_update_weights_hook() + + # Verify both signals were waited on in correct order + assert "update_signal" in update_signal.wait_calls, "Should wait for update signal" + assert "idle" in idle_event.wait_calls, "Should wait for idle" + + # Verify idle is cleared (blocks rollout during weight update) + assert policy._idle_as_infer.is_set() is False + + policy._post_update_weights_hook() + + # Verify idle is restored (allows rollout) + assert policy._idle_as_infer.is_set() is True + # Verify update signal is cleared + assert policy._update_weights_signal.is_set() is False + + +def test_update_weight_hooks_async_manage_events(): + """_pre/_post_update_weights_hook should manage async admission events.""" + + class _TrackingEvent: + """Event wrapper that tracks wait() calls.""" + + def __init__(self, name): + self._event = threading.Event() + self._event.set() + self._name = name + self.wait_calls = [] + + def wait(self): + self.wait_calls.append(self._name) + return self._event.wait() + + def set(self): + self._event.set() + + def clear(self): + self._event.clear() + + def is_set(self): + return self._event.is_set() + + policy = _DummyPolicy(_make_policy_config("async"), PolicyRole.EVAL) + policy.rollout_mode = "async" + + update_signal = _TrackingEvent("update_signal") + policy._update_weights_signal = update_signal + policy._weight_update_done = None + + policy._accept_new_requests = asyncio.Event() + policy._accept_new_requests.set() + policy.num_requests = 0 + policy._num_requests_lock = threading.Lock() + policy._inflight_zero_cv = threading.Condition(policy._num_requests_lock) + policy._loop = _FakeLoop() + + policy._pre_update_weights_hook() + + # Verify update signal was waited on + assert "update_signal" in update_signal.wait_calls, "Should wait for update signal" + + # Verify call_soon_threadsafe was used to clear accept_new_requests + assert len(policy._loop.calls) >= 1, "Should call call_soon_threadsafe" + + # Verify accept_new_requests is cleared (stops accepting new requests) + assert policy._accept_new_requests.is_set() is False + + policy._post_update_weights_hook() + + # Verify call_soon_threadsafe was used to set accept_new_requests + assert len(policy._loop.calls) >= 2, "Should call call_soon_threadsafe for set" + + # Verify accept_new_requests is restored (allows new requests) + assert policy._accept_new_requests.is_set() is True + # Verify update signal is cleared + assert policy._update_weights_signal.is_set() is False + + +def test_postprocess_sync_handles_remote_and_sets_idle(monkeypatch): + """_postprocess should convert outputs for remote and restore idle flag.""" + policy = _DummyPolicy(_make_policy_config("sync"), PolicyRole.EVAL) + policy._idle_as_infer = threading.Event() + policy._idle_as_infer.set() + + monkeypatch.setenv("RLIGHTNING_REMOTE_STORAGE", "1") + + env_ret = EnvRet(env_id="env-3", observation=np.array([1.0])) + env_cuda_called = {"called": False} + env_numpy_called = {"called": False} + + def _env_cuda(device=None): # noqa: ARG001 + env_cuda_called["called"] = True + return env_ret + + def _env_numpy(): + env_numpy_called["called"] = True + return env_ret + + env_ret.cuda = _env_cuda + env_ret.numpy = _env_numpy + + policy_resp = PolicyResponse(env_id="env-3", action=torch.tensor(1)) + resp_cuda_called = {"called": False} + resp_numpy_called = {"called": False} + + def _resp_cuda(device=None): # noqa: ARG001 + resp_cuda_called["called"] = True + return policy_resp + + def _resp_numpy(): + resp_numpy_called["called"] = True + return policy_resp + + policy_resp.cuda = _resp_cuda + policy_resp.numpy = _resp_numpy + + result = policy._postprocess(env_ret=env_ret, policy_resp=policy_resp) + + assert isinstance(result, tuple) + assert len(result) == 3 + assert env_cuda_called["called"] is True + assert resp_cuda_called["called"] is True + assert env_numpy_called["called"] is True + assert resp_numpy_called["called"] is True + assert isinstance(result[2], np.ndarray) + assert policy._idle_as_infer.is_set() is True + + +def test_check_idle_sync_allows_eval_and_async_asserts(): + """check_idle should work for sync eval and assert for async mode.""" + policy = _DummyPolicy(_make_policy_config("sync"), PolicyRole.EVAL) + policy._idle_as_infer = threading.Event() + policy._idle_as_infer.set() + + policy.check_idle() + + policy.rollout_mode = "async" + with pytest.raises(AssertionError): + policy.check_idle() + + +def test_get_num_requests_defaults_to_zero(): + """get_num_requests should return 0 when num_requests is None.""" + policy = _DummyPolicy(_make_policy_config("sync"), PolicyRole.EVAL) + policy.num_requests = None + + assert policy.get_num_requests() == 0 + + +def test_reset_training_state_rebuilds_and_clears(): + """reset_training_state should clear internal state and rebuild.""" + policy = _WarmupPolicy(_make_policy_config("sync"), PolicyRole.TRAIN) + policy.dataset = "data" + policy.optimizer = object() + policy.optimizer_steps = 5 + policy.model_list = [("x", nn.Linear(1, 1))] + policy._modules["rsl_rl/test"] = nn.Linear(1, 1) + + train_cfg = TrainConfig(max_epochs=1) + policy.reset_training_state(train_cfg, seed=42) + + assert policy.dataset is None + assert policy.optimizer is None + assert policy.optimizer_steps == 0 + assert policy.model_list == [] + assert "rsl_rl/test" not in policy._modules + assert policy.init_train_called is True + assert policy.setup_sampling_called is True + + +def test_save_weights_writes_state_dict(tmp_path): + """save_weights should write model state dicts to disk.""" + policy = _DummyPolicy(_make_policy_config("sync"), PolicyRole.TRAIN) + policy.construct_network() + policy._find_model() + + ckpt_path = tmp_path / "epoch_1" / "model.pt" + policy.save_checkpoint(ckpt_path) + + assert ckpt_path.exists() + state = torch.load(ckpt_path) + assert "linear" in state + + +def test_print_timing_summary_logs_and_resets(caplog): + """print_timing_summary should log timing entries and reset when requested.""" + policy = _DummyPolicy(_make_policy_config("sync"), PolicyRole.EVAL) + policy.timing_raw = {"rollout": {"count": 1, "total": 1.0, "avg": 1.0}} + + with caplog.at_level("DEBUG"): + policy.print_timing_summary(reset=True) + + assert "Policy" in caplog.text + assert "rollout" in caplog.text + assert policy.timing_raw == {} + + +def test_notify_update_weights_sets_signal(): + """notify_update_weights should set the update signal event.""" + policy = _DummyPolicy(_make_policy_config("sync"), PolicyRole.EVAL) + policy._update_weights_signal = threading.Event() + policy._weight_update_done = None + + policy.notify_update_weights() + + assert policy._update_weights_signal.is_set() is True + + +# ============================================================================ +# P0 - Core functionality tests for update_weights +# ============================================================================ + + +def test_update_weights_calls_hooks_and_buffer(monkeypatch): + """update_weights should call pre hook, buffer update, and post hook.""" + policy = _DummyPolicy(_make_policy_config("sync"), PolicyRole.EVAL) + policy.rollout_mode = "sync" + policy._update_weights_signal = threading.Event() + policy._update_weights_signal.set() + policy._idle_as_infer = threading.Event() + policy._idle_as_infer.set() + policy._weight_update_done = None + policy.timing_raw = {} + + call_order = [] + call_count = {"pre": 0} + + original_pre_hook = policy._pre_update_weights_hook + + def mock_pre_hook(): + call_count["pre"] += 1 + call_order.append("pre") + if call_count["pre"] >= 2: + raise KeyboardInterrupt # Terminate loop + original_pre_hook() + + def mock_update_buffer(): + call_order.append("buffer") + + original_post_hook = policy._post_update_weights_hook + + def mock_post_hook(): + call_order.append("post") + original_post_hook() + # Re-set the signal for next iteration + policy._update_weights_signal.set() + + monkeypatch.setattr(policy, "_pre_update_weights_hook", mock_pre_hook) + monkeypatch.setattr(policy, "update_weights_from_buffer", mock_update_buffer) + monkeypatch.setattr(policy, "_post_update_weights_hook", mock_post_hook) + + with pytest.raises(KeyboardInterrupt): + BasePolicy.update_weights(policy) + + assert call_order == ["pre", "buffer", "post", "pre"] + + +def test_update_weights_handles_exception_gracefully(monkeypatch): + """update_weights should catch exceptions from update_weights_from_buffer and continue.""" + policy = _DummyPolicy(_make_policy_config("sync"), PolicyRole.EVAL) + policy.rollout_mode = "sync" + policy._update_weights_signal = threading.Event() + policy._update_weights_signal.set() + policy._idle_as_infer = threading.Event() + policy._idle_as_infer.set() + policy._weight_update_done = None + policy.timing_raw = {} + + call_order = [] + call_count = {"pre": 0} + + original_pre_hook = policy._pre_update_weights_hook + + def mock_pre_hook(): + call_count["pre"] += 1 + call_order.append("pre") + if call_count["pre"] >= 2: + raise KeyboardInterrupt # Terminate loop after second iteration + original_pre_hook() + + def mock_update_buffer_raises(): + call_order.append("buffer_error") + raise ValueError("test buffer error") + + original_post_hook = policy._post_update_weights_hook + + def mock_post_hook(): + call_order.append("post") + original_post_hook() + policy._update_weights_signal.set() + + monkeypatch.setattr(policy, "_pre_update_weights_hook", mock_pre_hook) + monkeypatch.setattr(policy, "update_weights_from_buffer", mock_update_buffer_raises) + monkeypatch.setattr(policy, "_post_update_weights_hook", mock_post_hook) + + with pytest.raises(KeyboardInterrupt): + BasePolicy.update_weights(policy) + + # Exception was caught, post hook still called (finally block) + assert call_order == ["pre", "buffer_error", "post", "pre"] + + +# ============================================================================ +# P1 - Exception handling tests for rollout and postprocess +# ============================================================================ + + +def test_rollout_logs_and_reraises_exception(): + """_rollout should log and re-raise exceptions from rollout_step.""" + policy = _FailingRolloutPolicy(_make_policy_config("sync"), PolicyRole.EVAL) + policy._idle_as_infer = threading.Event() + policy._idle_as_infer.set() + + env_ret = EnvRet(env_id="env-fail", observation=np.array([1.0])) + + with pytest.raises(RuntimeError, match="test rollout error"): + policy._rollout(env_ret) + + +def test_rollout_async_logs_and_reraises_exception(): + """_rollout_async should log and re-raise exceptions from rollout_step.""" + + class _FailingAsyncRolloutPolicy(_DummyPolicy): + """Async policy that raises in rollout_step.""" + + async def rollout_step(self, env_ret: EnvRet, **kwargs: Any) -> PolicyResponse: # noqa: ARG002 + raise RuntimeError("test async rollout error") + + async def _run(): + policy = _FailingAsyncRolloutPolicy(_make_policy_config("async"), PolicyRole.EVAL) + policy.rollout_mode = "async" + policy._accept_new_requests = asyncio.Event() + policy._accept_new_requests.set() + policy.num_requests = 0 + policy._num_requests_lock = threading.Lock() + policy._inflight_zero_cv = threading.Condition(policy._num_requests_lock) + + env_ret = EnvRet(env_id="env-fail-async", observation=np.array([0.0])) + await policy._rollout_async(env_ret) + + with pytest.raises(RuntimeError, match="test async rollout error"): + asyncio.run(_run()) + + +def test_postprocess_logs_and_reraises_exception(): + """_postprocess should log and re-raise exceptions from postprocess.""" + policy = _FailingPostprocessPolicy(_make_policy_config("sync"), PolicyRole.EVAL) + policy._idle_as_infer = threading.Event() + policy._idle_as_infer.set() + + env_ret = EnvRet(env_id="env-pp-fail", observation=np.array([1.0])) + policy_resp = PolicyResponse(env_id="env-pp-fail", action=torch.tensor(1)) + + with pytest.raises(RuntimeError, match="test postprocess error"): + policy._postprocess(env_ret=env_ret, policy_resp=policy_resp) + + +# ============================================================================ +# P2 - Enhanced coverage tests +# ============================================================================ + + +def test_pre_rollout_hook_records_timing_when_debug(monkeypatch): + """_pre_rollout_hook should record timing when DEBUG flag is set.""" + policy = _DummyPolicy(_make_policy_config("sync"), PolicyRole.EVAL) + policy._idle_as_infer = threading.Event() + policy._idle_as_infer.set() + policy.timing_raw = {} + + # Enable DEBUG flag + monkeypatch.setenv("RLIGHTNING_DEBUG", "1") + + env_ret = EnvRet(env_id="env-debug", observation=np.array([1.0])) + + record_called = {"called": False} + + def mock_record_timing(name, _value, _timing_dict, level="debug"): # noqa: ARG001 + record_called["called"] = True + assert name == "transition_env_to_policy" + + monkeypatch.setattr("rlightning.policy.base_policy.profiler.record_timing", mock_record_timing) + + policy._pre_rollout_hook(env_ret) + + assert record_called["called"] is True + + +def test_pre_rollout_hook_async_records_timing_when_debug(monkeypatch): + """_pre_rollout_hook_async should record timing when DEBUG flag is set.""" + + async def _run(): + policy = _AsyncDummyPolicy(_make_policy_config("async"), PolicyRole.EVAL) + policy.rollout_mode = "async" + policy._accept_new_requests = asyncio.Event() + policy._accept_new_requests.set() + policy.num_requests = 0 + policy._num_requests_lock = threading.Lock() + policy.timing_raw = {} + + env_ret = EnvRet(env_id="env-debug-async", observation=np.array([0.0])) + env_ret.ts_env_sent_ns = 0 + + record_called = {"called": False} + + def mock_record_timing(name, _value, _timing_dict, level="debug"): # noqa: ARG001 + record_called["called"] = True + assert name == "transition_env_to_policy" + + monkeypatch.setattr("rlightning.policy.base_policy.profiler.record_timing", mock_record_timing) + + await policy._pre_rollout_hook_async(env_ret) + + return record_called["called"] + + # Enable DEBUG flag + monkeypatch.setenv("RLIGHTNING_DEBUG", "1") + + record_called = asyncio.run(_run()) + assert record_called is True + + +def test_get_num_requests_returns_actual_count(): + """get_num_requests should return actual count when num_requests is set.""" + policy = _DummyPolicy(_make_policy_config("async"), PolicyRole.EVAL) + policy.num_requests = 5 + + assert policy.get_num_requests() == 5 + + +def test_post_postprocess_hook_handles_list_result(monkeypatch): + """_post_postprocess_hook should convert list results when remote.""" + policy = _DummyPolicy(_make_policy_config("sync"), PolicyRole.EVAL) + policy._idle_as_infer = threading.Event() + + # Enable remote flag + monkeypatch.setenv("RLIGHTNING_REMOTE_STORAGE", "1") + + numpy_calls = {"count": 0} + + def make_mock_numpy(original): + def _numpy(): + numpy_calls["count"] += 1 + return original + + return _numpy + + env_ret1 = EnvRet(env_id="env-list-1", observation=np.array([1.0])) + env_ret1.numpy = make_mock_numpy(env_ret1) + + env_ret2 = EnvRet(env_id="env-list-2", observation=np.array([2.0])) + env_ret2.numpy = make_mock_numpy(env_ret2) + + result = [env_ret1, env_ret2] + + converted = policy._post_postprocess_hook(result) + + assert isinstance(converted, list) + assert len(converted) == 2 + assert numpy_calls["count"] == 2 + assert policy._idle_as_infer.is_set() is True + + +def test_check_idle_allows_train_role(): + """check_idle should pass for TRAIN role without waiting.""" + policy = _DummyPolicy(_make_policy_config("sync"), PolicyRole.TRAIN) + policy._idle_as_infer = threading.Event() + # Don't set the event - would block if EVAL role waited + + # Should not raise and not block + policy.check_idle() + + +def test_post_rollout_hook_skips_cpu_when_not_remote(monkeypatch): + """_post_rollout_hook should not call numpy() when not in remote mode.""" + policy = _DummyPolicy(_make_policy_config("sync"), PolicyRole.EVAL) + policy._idle_as_infer = threading.Event() + + # Ensure remote flags are NOT set + monkeypatch.setenv("RLIGHTNING_REMOTE_STORAGE", "0") + monkeypatch.setenv("RLIGHTNING_REMOTE_ENV", "0") + + numpy_called = {"called": False} + + policy_resp = PolicyResponse(env_id="env-local", action=torch.tensor(1)) + + def _numpy(): + numpy_called["called"] = True + return policy_resp + + policy_resp.numpy = _numpy + + result = policy._post_rollout_hook(policy_resp) + + assert numpy_called["called"] is False + assert result is policy_resp + assert policy._idle_as_infer.is_set() is True + + +def test_save_weights_unwraps_ddp_module(tmp_path, monkeypatch): + """save_weights should unwrap DDP modules and save inner module state.""" + from torch.nn.parallel import DistributedDataParallel as DDP + + policy = _DummyPolicy(_make_policy_config("sync"), PolicyRole.TRAIN) + policy.construct_network() + + # Create a mock DDP wrapper that passes isinstance check + class MockDDP(nn.Module): + def __init__(self, module): + super().__init__() + self.module = module + + original_linear = policy.linear + wrapped_linear = MockDDP(original_linear) + + # Patch isinstance to recognize MockDDP as DDP + original_isinstance = builtins.isinstance + + def patched_isinstance(obj, classinfo): + if classinfo is DDP and type(obj).__name__ == "MockDDP": + return True + return original_isinstance(obj, classinfo) + + monkeypatch.setattr(builtins, "isinstance", patched_isinstance) + + policy.model_list = [("linear", wrapped_linear)] + + ckpt_path = tmp_path / "epoch_2" / "model.pt" + policy.save_checkpoint(ckpt_path) + + assert ckpt_path.exists() + state = torch.load(ckpt_path) + assert "linear" in state + # Verify the state dict came from the inner module + assert "weight" in state["linear"] + assert "bias" in state["linear"] + + +def test_save_checkpoint_without_model_attribute(tmp_path): + """save_checkpoint should work for policies that only populate model_list.""" + policy = _DummyPolicy(_make_policy_config("sync"), PolicyRole.TRAIN) + policy.construct_network() + policy._find_model() + if hasattr(policy, "model"): + delattr(policy, "model") + + ckpt_path = tmp_path / "epoch_3" / "model.pt" + policy.save_checkpoint(ckpt_path) + + state = torch.load(ckpt_path, map_location="cpu", weights_only=False) + assert "linear" in state + assert "weight" in state["linear"] diff --git a/tests/unit/test_base_policy_checkpoint.py b/tests/unit/test_base_policy_checkpoint.py new file mode 100644 index 0000000..1525f5a --- /dev/null +++ b/tests/unit/test_base_policy_checkpoint.py @@ -0,0 +1,121 @@ +"""Focused tests for BasePolicy checkpoint serialization behavior.""" + +from __future__ import annotations + +from typing import Any, Dict + +import torch +from torch import nn +from torch.distributed.utils import _free_storage + +from rlightning.policy.base_policy import BasePolicy, PolicyRole, clone_checkpoint_value +from rlightning.types import EnvRet, PolicyResponse +from rlightning.utils.config import PolicyConfig + + +class _CheckpointPolicy(BasePolicy): + """Concrete policy used to exercise checkpoint save paths.""" + + def construct_network(self, env_meta: Any = None, *args: Any, **kwargs: Any) -> None: + self.linear = nn.Linear(2, 2) + + def setup_optimizer(self, optim_cfg: Any) -> None: + self.optimizer = torch.optim.SGD(self.parameters(), lr=0.1) + + def rollout_step(self, env_ret: EnvRet, **kwargs: Any) -> PolicyResponse: + return PolicyResponse(env_id=env_ret.env_id, action=torch.tensor(1)) + + def postprocess(self, env_ret: EnvRet | None = None, policy_resp: PolicyResponse | None = None) -> Any: + return env_ret, policy_resp + + def update_dataset(self, data: Any) -> None: + self.dataset = data + + def train(self, *args: Any, **kwargs: Any) -> Any: + return None + + def get_trainable_parameters(self) -> Dict[str, Dict[str, torch.Tensor]]: + return {name: module.state_dict() for name, module in self.model_list} + + def load_state_dict(self, state_dict: Dict[str, torch.Tensor], *args: Any, **kwargs: Any) -> None: + self.loaded_state = state_dict + + +def _make_policy() -> _CheckpointPolicy: + config = PolicyConfig.from_dict({"type": "CheckpointPolicy", "rollout_mode": "sync"}) + policy = _CheckpointPolicy(config, PolicyRole.TRAIN) + policy.construct_network() + policy._find_model() + policy.model = policy.linear + return policy + + +def test_clone_checkpoint_value_produces_loadable_cpu_tensors(tmp_path): + """clone_checkpoint_value should detach tensor payloads into standalone CPU storages.""" + source = { + "view": torch.arange(8, dtype=torch.float32).reshape(2, 4)[:, :2], + "nested": (torch.arange(4, dtype=torch.bfloat16).reshape(2, 2),), + } + + cloned = clone_checkpoint_value(source) + + assert cloned["view"].device.type == "cpu" + assert cloned["nested"][0].device.type == "cpu" + assert cloned["view"].data_ptr() != source["view"].data_ptr() + + checkpoint_path = tmp_path / "checkpoint.pt" + torch.save(cloned, checkpoint_path) + loaded = torch.load(checkpoint_path, map_location="cpu", weights_only=False) + + torch.testing.assert_close(loaded["view"], source["view"]) + torch.testing.assert_close(loaded["nested"][0], source["nested"][0]) + + +def test_save_checkpoint_round_trips_offloaded_parameters(tmp_path, monkeypatch): + """save_checkpoint should temporarily reload an offloaded model and restore its state.""" + policy = _make_policy() + original_state = {name: tensor.detach().cpu().clone() for name, tensor in policy.linear.state_dict().items()} + monkeypatch.setattr(policy, "clear_memory", lambda sync=False: None) + monkeypatch.setattr( + "rlightning.weights.weight_buffer_mixin.profiler.log_gpu_memory_usage", + lambda *args, **kwargs: None, + ) + + for name, param in policy.linear.named_parameters(): + policy.cpu_param_backup[name] = (param.data.detach().cpu().clone(), param.data.size()) + _free_storage(param.data) + policy._model_params_offloaded = True + + for param in policy.linear.parameters(): + assert param.data.storage().size() == 0 + + checkpoint_path = tmp_path / "epoch_1" / "model.pt" + policy.save_checkpoint(checkpoint_path) + + assert checkpoint_path.exists() + state = torch.load(checkpoint_path, map_location="cpu", weights_only=False) + assert "linear" in state + for name, tensor in original_state.items(): + torch.testing.assert_close(state["linear"][name], tensor) + + for param in policy.linear.parameters(): + assert param.data.storage().size() == 0 + + +def test_save_checkpoint_does_not_treat_stale_cpu_backup_as_offloaded(tmp_path): + """Existing cpu backups should not affect checkpointing when the model is not offloaded.""" + policy = _make_policy() + + mutated_weight = torch.full_like(policy.linear.weight, 7.0) + backup_weight = torch.full_like(policy.linear.weight, -3.0) + + with torch.no_grad(): + policy.linear.weight.copy_(mutated_weight) + policy.cpu_param_backup["weight"] = (backup_weight.clone(), policy.linear.weight.data.size()) + policy._model_params_offloaded = False + + checkpoint_path = tmp_path / "epoch_2" / "model.pt" + policy.save_checkpoint(checkpoint_path) + + state = torch.load(checkpoint_path, map_location="cpu", weights_only=False) + torch.testing.assert_close(state["linear"]["weight"], mutated_weight.cpu()) diff --git a/tests/unit/test_buffer_utils.py b/tests/unit/test_buffer_utils.py new file mode 100644 index 0000000..7d36406 --- /dev/null +++ b/tests/unit/test_buffer_utils.py @@ -0,0 +1,163 @@ +import torch +import pytest + +from rlightning.buffer.utils.utils import ( + default_compute_gae, + default_env_ret_preprocess_fn, + default_postprocess_fn, + default_policy_resp_preprocess_fn, + default_preprocess_fn, +) +from rlightning.types import EnvRet, PolicyResponse + + +def test_default_env_ret_preprocess_fn_applies_obs_and_reward_preprocessors(): + env_ret = EnvRet( + env_id="env-1", + observation=[1.0, 2.0], + last_reward=3.0, + last_terminated=False, + last_truncated=False, + info={"k": 1}, + ) + + def obs_preprocessor(x): + return torch.tensor(x) + 1 + + def reward_preprocessor(x): + return torch.tensor(x) * 2 + + transition = default_env_ret_preprocess_fn({}, env_ret, obs_preprocessor, reward_preprocessor) + + expected_obs = torch.tensor([2.0, 3.0], dtype=transition["observation"].dtype) + expected_reward = torch.tensor(6.0, dtype=transition["last_reward"].dtype) + + assert torch.allclose(transition["observation"], expected_obs) + assert torch.allclose(transition["last_reward"], expected_reward) + assert transition["last_terminated"] is False + assert transition["last_truncated"] is False + assert transition["info"] == {"k": 1} + + +def test_default_env_ret_preprocess_fn_rejects_non_env_return_input(): + with pytest.raises(TypeError): + default_env_ret_preprocess_fn({}, "not-env-ret", lambda x: x, lambda x: x) + + +def test_default_policy_resp_preprocess_fn_merges_policy_fields(): + transition = default_policy_resp_preprocess_fn({}, PolicyResponse(env_id="env-1", action=torch.tensor(1), log_prob=0.5)) + + assert torch.equal(transition["action"], torch.tensor(1)) + assert transition["log_prob"] == 0.5 + + +def test_default_policy_resp_preprocess_fn_rejects_non_policy_response_input(): + with pytest.raises(TypeError): + default_policy_resp_preprocess_fn({}, "not-policy") + + +def test_default_preprocess_fn_requires_at_least_one_input(): + with pytest.raises(ValueError, match="At least one of env_ret or policy_resp"): + default_preprocess_fn({}) + + +def test_default_preprocess_fn_rejects_mismatched_env_ids(): + env_ret = EnvRet(env_id="env-0", observation=[1, 2], last_reward=1.0) + policy_resp = PolicyResponse(env_id="env-1", action=0) + + with pytest.raises(ValueError, match="Mismatched env_id"): + default_preprocess_fn({}, env_ret=env_ret, policy_resp=policy_resp) + + +def test_default_preprocess_fn_merges_env_return_and_policy_response(): + env_ret = EnvRet( + env_id="env-0", + observation={"state": [1.0, 2.0]}, + last_reward=1.5, + last_terminated=False, + last_truncated=False, + info={"episode": 3}, + ) + policy_resp = PolicyResponse(env_id="env-0", action=2, log_prob=0.3) + + transition = default_preprocess_fn({}, env_ret=env_ret, policy_resp=policy_resp) + + assert transition["observation"] == {"state": [1.0, 2.0]} + assert transition["last_reward"] == 1.5 + assert transition["info"] == {"episode": 3} + assert transition["action"] == 2 + assert transition["log_prob"] == 0.3 + + +def test_default_postprocess_fn_builds_training_batch_from_episode_buffer(): + episode_buffer = { + "observation": [ + torch.tensor([1.0, 2.0]), + torch.tensor([3.0, 4.0]), + torch.tensor([5.0, 6.0]), + ], + "last_reward": [0.0, 1.0, 2.0], + "last_terminated": [False, False, True], + "last_truncated": [False, False, False], + "info": [{"step": 0}, {"step": 1}, {"step": 2}], + "action": [ + torch.tensor([0.1]), + torch.tensor([0.2]), + torch.tensor([0.3]), + ], + "log_prob": [torch.tensor(0.1), torch.tensor(0.2), torch.tensor(0.3)], + } + + data = default_postprocess_fn(episode_buffer) + + assert "info" not in data + assert torch.equal(data["observation"], torch.tensor([[1.0, 2.0], [3.0, 4.0]])) + assert torch.equal(data["next_observation"], torch.tensor([[3.0, 4.0], [5.0, 6.0]])) + assert torch.equal(data["reward"], torch.tensor([1.0, 2.0])) + assert torch.equal(data["terminated"], torch.tensor([False, True])) + assert torch.equal(data["truncated"], torch.tensor([False, False])) + assert torch.equal(data["action"], torch.tensor([[0.1], [0.2]])) + assert torch.allclose(data["log_prob"], torch.tensor([0.1, 0.2])) + + +def test_default_compute_gae_respects_terminal_boundaries(): + rewards = torch.tensor([[1.0], [2.0], [3.0]]) + values = torch.tensor([[0.5], [0.5], [0.5]]) + next_values = torch.tensor([[0.5], [0.5], [0.0]]) + dones = torch.tensor([[0.0], [0.0], [1.0]]) + + advantages, returns = default_compute_gae( + rewards=rewards, + values=values, + next_values=next_values, + dones=dones, + gamma=0.9, + lam=0.95, + normalize_adv=False, + ) + + expected_advantages = torch.tensor([[4.4448], [4.0875], [2.5000]]) + expected_returns = torch.tensor([[4.9448], [4.5875], [3.0000]]) + + assert torch.allclose(advantages, expected_advantages, atol=1e-4) + assert torch.allclose(returns, expected_returns, atol=1e-4) + + +def test_default_compute_gae_normalizes_advantages_when_requested(): + rewards = torch.tensor([[1.0], [2.0], [3.0]]) + values = torch.tensor([[0.5], [0.5], [0.5]]) + next_values = torch.tensor([[0.5], [0.5], [0.5]]) + dones = torch.tensor([[0.0], [0.0], [0.0]]) + + advantages, _ = default_compute_gae( + rewards=rewards, + values=values, + next_values=next_values, + dones=dones, + gamma=0.9, + lam=0.95, + normalize_adv=True, + ) + + assert torch.allclose(advantages.mean(), torch.tensor(0.0), atol=1e-6) + assert torch.allclose(advantages.std(), torch.tensor(1.0), atol=1e-6) diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py new file mode 100644 index 0000000..d86b83d --- /dev/null +++ b/tests/unit/test_config.py @@ -0,0 +1,131 @@ +import pytest +from omegaconf import OmegaConf + +from rlightning.utils.config import ( + BufferConfig, + Config, + EnvConfig, + LogConfig, + MainConfig, + PolicyConfig, + TrainConfig, + WeightBufferConfig, + validate_config_for_placement, +) +import rlightning.utils.placement as placement_module + + +def test_buffer_config_sets_default_sampler_by_buffer_type(): + rollout_cfg = BufferConfig(type="RolloutBuffer", capacity=8) + replay_cfg = BufferConfig(type="ReplayBuffer", capacity=8) + + assert rollout_cfg.sampler.type == "all" + assert replay_cfg.sampler.type == "uniform" + + +def test_weight_buffer_config_rejects_shared_strategy_for_non_cpu_weight_buffer(): + with pytest.raises(ValueError, match="Shared buffer strategy"): + WeightBufferConfig(type="WeightBuffer", buffer_strategy="Shared") + + +def test_env_config_rejects_vector_env_for_unsupported_backend(): + with pytest.raises(ValueError, match="Vectorized environments are only supported"): + EnvConfig( + name="bad-env", + backend="mujoco", + task="cartpole", + num_envs=2, + ) + + +def test_env_config_sets_maniskill_control_mode(): + cfg = EnvConfig( + name="maniskill-env", + backend="maniskill", + task="pick_cube", + init_params={}, + policy_setup="widowx", + ) + + assert cfg.init_params.control_mode == "arm_pd_ee_target_delta_pose_align2_gripper_pd_joint_pos" + + +def test_main_config_from_dict_formats_validation_errors(make_main_config_dict): + invalid_config = make_main_config_dict() + del invalid_config["train"]["max_epochs"] + + with pytest.raises(ValueError, match="field 'train.max_epochs'"): + MainConfig.from_dict(invalid_config) + + +def test_main_config_from_omegaconf_promotes_nested_config_types(make_main_config_dict): + cfg = MainConfig.from_omegaconf(OmegaConf.create(make_main_config_dict())) + + assert isinstance(cfg.env, list) + assert isinstance(cfg.env[0], EnvConfig) + assert isinstance(cfg.buffer, BufferConfig) + assert isinstance(cfg.policy, PolicyConfig) + assert isinstance(cfg.train, TrainConfig) + assert isinstance(cfg.log, LogConfig) + + +def test_main_config_converts_single_env_config_to_list(make_main_config_dict): + config_dict = make_main_config_dict() + config_dict["env"] = config_dict["env"][0] + + config = MainConfig.from_dict(config_dict) + + assert isinstance(config.env, list) + assert len(config.env) == 1 + assert isinstance(config.env[0], EnvConfig) + + +def test_config_allows_extra_fields_and_recursively_wraps_extra_dicts(make_main_config_dict): + config_dict = make_main_config_dict() + config_dict["extra_level_1"] = {"level_2": {"param": 42}} + + config = MainConfig.from_dict(config_dict) + + assert isinstance(config.extra_level_1, Config) + assert isinstance(config.extra_level_1.level_2, Config) + assert config.extra_level_1.level_2.param == 42 + + +def test_config_get_and_getitem_return_field_values(): + cfg = Config.from_dict({"alpha": 1, "nested": {"beta": 2}}) + + assert cfg.get("alpha") == 1 + assert cfg["alpha"] == 1 + assert cfg.get("missing", "fallback") == "fallback" + assert isinstance(cfg.nested, Config) + assert cfg.nested.beta == 2 + + +def test_main_config_to_dict_and_yaml_include_nested_fields(make_main_config_dict): + config = MainConfig.from_dict(make_main_config_dict()) + + as_dict = config.to_dict() + as_yaml = config.to_yaml() + + assert as_dict["buffer"]["type"] == "RolloutBuffer" + assert as_dict["policy"]["type"] == "SimplePPOPolicy" + assert "buffer:" in as_yaml + assert "policy:" in as_yaml + + +def test_validate_config_for_placement_applies_colocate_overrides(make_main_config_dict, monkeypatch): + class FakeResourceManager: + def get_placement_strategy(self) -> str: + return "colocate" + + monkeypatch.setattr(placement_module, "get_global_resource_manager", lambda: FakeResourceManager()) + + config = MainConfig.from_dict(make_main_config_dict()) + validated = validate_config_for_placement(config) + + assert validated.cluster.train_each_gpu_num == 0.1 + assert validated.cluster.eval_each_gpu_num == 0.1 + assert validated.cluster.is_colocated is True + assert validated.policy.weight_buffer.buffer_strategy == "None" + assert validated.env[0].num_gpus == 0.1 + assert validated.env[0].num_cpus == 1 diff --git a/tests/unit/test_env_group_stats.py b/tests/unit/test_env_group_stats.py new file mode 100644 index 0000000..48a035b --- /dev/null +++ b/tests/unit/test_env_group_stats.py @@ -0,0 +1,47 @@ +"""Unit tests for EnvGroup.get_env_stats.""" + +import pytest + +from rlightning.env.env_group import EnvGroup + +class _DummyEnv: + def __init__(self, stats): + self._stats = stats + self.reset_flags = [] + self.finish_rollout_called = 0 + + def get_env_stats(self, reset=False): + self.reset_flags.append(reset) + return self._stats + + def finish_rollout(self): + self.finish_rollout_called += 1 + + +class _DummySubmitter: + def submit(self, fn, *args, _block: bool = False, **kwargs): + return fn(*args, **kwargs) + + +def _make_env_group(env_list, env_servers): + env_group = EnvGroup.__new__(EnvGroup) + env_group.env_list = env_list + env_group.env_servers = env_servers + env_group._task_submitter = _DummySubmitter() + return env_group + + +def test_get_env_stats_aggregates_means_across_envs_and_servers(): + env_local = _DummyEnv(stats={"reward": [3.0, 2], "success": [1.0, 1]}) + env_server = _DummyEnv(stats={"reward": [5.0, 2], "fail": [2.0, 4]}) + env_group = _make_env_group([env_local], [env_server]) + + stats = env_group.get_env_stats(reset=True) + + assert stats["reward"] == pytest.approx(2.0) # (3 + 5) / (2 + 2) + assert stats["success"] == pytest.approx(1.0) # 1 / 1 + assert stats["fail"] == pytest.approx(0.5) # 2 / 4 + assert env_local.reset_flags == [True] + assert env_server.reset_flags == [True] + assert env_local.finish_rollout_called == 1 + assert env_server.finish_rollout_called == 1 diff --git a/tests/unit/test_placement_manager.py b/tests/unit/test_placement_manager.py new file mode 100644 index 0000000..29e18a9 --- /dev/null +++ b/tests/unit/test_placement_manager.py @@ -0,0 +1,342 @@ +from __future__ import annotations + +from pathlib import Path + +import ray +import pytest +import yaml + +from rlightning.utils.config import ClusterConfig +from rlightning.utils.placement import ComponentScheduling, GlobalResourceManager, ResourcePoolPlanner, Scheduling +import rlightning.utils.placement.placement_manager as placement_manager_module +import rlightning.utils.placement.placement_strategies as placement_strategies_module +import rlightning.utils.placement.resource_pool as resource_pool_module + + +def create_mock_cluster_resources( + num_nodes: int = 2, + gpus_per_node: int = 8, + cpus_per_node: int = 64, +): + node_id_to_resources = {} + for i in range(num_nodes): + node_id = f"node_{i}" + node_id_to_resources[node_id] = { + "node_id": node_id, + "ip": f"192.168.1.{i + 1}", + "CPU": cpus_per_node, + "GPU": gpus_per_node, + } + return { + "node_id_to_resources": node_id_to_resources, + "total_cpus": num_nodes * cpus_per_node, + "total_gpus": num_nodes * gpus_per_node, + } + + +def create_mock_scheduling( + train_workers: int = 0, + train_gpus: float = 1.0, + eval_workers: int = 0, + eval_gpus: float = 1.0, + buffer_workers: int | str = 0, + buffer_gpus: float = 0.0, + env_workers: int = 0, + env_gpus: float = 1.0, +) -> ComponentScheduling: + return ComponentScheduling( + train_worker=Scheduling(worker_num=train_workers, num_cpus=1, num_gpus=train_gpus), + eval_worker=Scheduling(worker_num=eval_workers, num_cpus=1, num_gpus=eval_gpus), + buffer_worker=Scheduling(worker_num=buffer_workers, num_cpus=1, num_gpus=buffer_gpus), + env_worker=[Scheduling(worker_num=env_workers, num_cpus=1, num_gpus=env_gpus)], + ) + + +def create_mock_cluster_config( + strategy: str = "default", + env_strategy: str = "default", + mode: str = "auto", + resource_pool: list[dict] | None = None, + train_worker_num: int = 4, + eval_worker_num: int = 2, + buffer_worker_num: int | str = 1, +) -> ClusterConfig: + cluster_dict = { + "placement": { + "strategy": strategy, + "env_strategy": env_strategy, + "mode": mode, + }, + "train_worker_num": train_worker_num, + "eval_worker_num": eval_worker_num, + "buffer_worker_num": buffer_worker_num, + } + if resource_pool is not None: + cluster_dict["resource_pool"] = resource_pool + return ClusterConfig.from_dict(cluster_dict) + + +class _FakePlacementGroup: + def __init__(self, bundles, name, strategy, node_id): + self.bundles = bundles + self.name = name + self.strategy = strategy + self.node_id = node_id + self.id = f"{name}-id" + + def ready(self): + return True + + +def patch_cluster_sources(monkeypatch, cluster_info): + monkeypatch.setattr(placement_manager_module, "get_cluster_resources", lambda: cluster_info) + monkeypatch.setattr(resource_pool_module, "get_cluster_resources", lambda: cluster_info) + monkeypatch.setattr(placement_strategies_module, "get_cluster_resources", lambda: cluster_info) + + +def patch_ray_runtime(monkeypatch): + created = {} + + def _placement_group(bundles, name, strategy, _soft_target_node_id=None): + pg = _FakePlacementGroup(bundles, name, strategy, _soft_target_node_id) + created[name] = pg + return pg + + monkeypatch.setattr(ray.util, "placement_group", _placement_group) + monkeypatch.setattr(ray, "get", lambda _ready: None) + return created + + +@pytest.fixture +def mock_cluster_2_nodes_8_gpus(): + return create_mock_cluster_resources(num_nodes=2, gpus_per_node=8) + + +@pytest.fixture +def mock_cluster_1_node_8_gpus(): + return create_mock_cluster_resources(num_nodes=1, gpus_per_node=8) + + +@pytest.fixture +def basic_scheduling(): + return create_mock_scheduling( + train_workers=4, + train_gpus=1.0, + eval_workers=2, + eval_gpus=1.0, + buffer_workers=1, + env_workers=4, + env_gpus=0.5, + ) + + +@pytest.fixture +def medium_scheduling(): + return create_mock_scheduling( + train_workers=8, + train_gpus=1.0, + eval_workers=4, + eval_gpus=1.0, + buffer_workers=1, + env_workers=8, + env_gpus=0.5, + ) + + +@pytest.fixture +def large_scheduling(): + return create_mock_scheduling( + train_workers=8, + train_gpus=1.0, + eval_workers=8, + eval_gpus=1.0, + buffer_workers=1, + env_workers=16, + env_gpus=0.5, + ) + + +def test_resource_pool_planner_discovers_cluster_resources(mock_cluster_2_nodes_8_gpus, basic_scheduling): + planner = ResourcePoolPlanner(scheduling=basic_scheduling, cluster_info=mock_cluster_2_nodes_8_gpus) + + nodes = planner.discover_cluster_resources() + + assert set(nodes) == {"node_0", "node_1"} + assert nodes["node_0"].total_gpus == 8 + assert nodes["node_1"].total_cpus == 64 + + +def test_resource_pool_planner_validates_disaggregate_scheduling( + mock_cluster_2_nodes_8_gpus, basic_scheduling +): + planner = ResourcePoolPlanner(scheduling=basic_scheduling, cluster_info=mock_cluster_2_nodes_8_gpus) + + is_valid, error_msg = planner.validate_scheduling(strategy="disaggregate") + + assert is_valid is True + assert error_msg == "" + assert planner.resource_summary == { + "train_pool_required_gpus": 4.0, + "rollout_pool_required_gpus": 4.0, + "env_required_gpus": 2.0, + "eval_required_gpus": 2.0, + } + + +def test_resource_pool_planner_rejects_insufficient_gpus(mock_cluster_1_node_8_gpus, large_scheduling): + planner = ResourcePoolPlanner(scheduling=large_scheduling, cluster_info=mock_cluster_1_node_8_gpus) + + is_valid, error_msg = planner.validate_scheduling(strategy="disaggregate") + + assert is_valid is False + assert "Insufficient GPUs" in error_msg + + +def test_resource_pool_planner_plans_disaggregate_pools(mock_cluster_2_nodes_8_gpus, medium_scheduling): + planner = ResourcePoolPlanner(scheduling=medium_scheduling, cluster_info=mock_cluster_2_nodes_8_gpus) + + pools = planner.plan_resource_pools(strategy="disaggregate") + + assert set(pools) == {"train_pool", "rollout_pool"} + assert set(pools["train_pool"].component_types) == {"buffer", "train"} + assert pools["train_pool"].get_component_indices("train") == "0-7" + assert set(pools["rollout_pool"].component_types) == {"env", "eval"} + assert pools["rollout_pool"].get_component_indices("eval") == "0-3" + assert pools["rollout_pool"].get_component_indices("env") == "4-7" + assert planner.get_pool_for_component("train").name == "train_pool" + assert planner.get_pool_for_component("eval").name == "rollout_pool" + + +def test_resource_pool_planner_loads_manual_resource_pools(mock_cluster_2_nodes_8_gpus, basic_scheduling): + planner = ResourcePoolPlanner(scheduling=basic_scheduling, cluster_info=mock_cluster_2_nodes_8_gpus) + planner.discover_cluster_resources() + + pools = planner.load_manual_resource_pools( + [ + { + "name": "train_pool", + "num_node": 1, + "num_gpus": 8, + "train": "0-7", + }, + { + "name": "rollout_pool", + "num_node": 1, + "num_gpus": 8, + "eval": "0-3", + "env": "4-7", + }, + ] + ) + + assert set(pools) == {"train_pool", "rollout_pool"} + assert set(pools["train_pool"].component_types) == {"buffer", "train"} + assert set(pools["rollout_pool"].component_types) == {"env", "eval"} + + +def test_resource_pool_planner_to_yaml_and_summary(mock_cluster_2_nodes_8_gpus, medium_scheduling): + planner = ResourcePoolPlanner(scheduling=medium_scheduling, cluster_info=mock_cluster_2_nodes_8_gpus) + planner.plan_resource_pools(strategy="disaggregate") + + yaml_config = planner.to_yaml_config() + summary = planner.summary() + + assert len(yaml_config) == 2 + assert set(yaml_config[0].keys()) == {"name", "num_node", "num_gpus", "train"} + assert set(yaml_config[1].keys()) == {"name", "num_node", "num_gpus", "eval", "env"} + assert set(summary) == {"cluster", "pools", "yaml_config"} + assert summary["yaml_config"] == yaml_config + + +def test_global_resource_manager_is_singleton(): + manager1 = GlobalResourceManager.get_instance() + manager2 = GlobalResourceManager.get_instance() + + assert manager1 is manager2 + assert manager1.is_initialized is False + + +def test_global_resource_manager_requires_initialize_for_runtime_methods(): + manager = GlobalResourceManager.get_instance() + + with pytest.raises(RuntimeError, match="not initialized"): + manager.get_scheduling_strategy("train", 0) + + with pytest.raises(RuntimeError, match="not initialized"): + manager.get_storage_to_train_workers() + + with pytest.raises(RuntimeError, match="not initialized"): + manager.save_yaml_config("/tmp") + + assert manager.get_placement_config() is None + assert manager.get_scheduling() is None + assert manager.get_pool_for_component("train") is None + + +def test_global_resource_manager_initializes_disaggregate_strategy( + monkeypatch, mock_cluster_2_nodes_8_gpus, medium_scheduling +): + patch_cluster_sources(monkeypatch, mock_cluster_2_nodes_8_gpus) + created = patch_ray_runtime(monkeypatch) + + manager = GlobalResourceManager.get_instance() + manager.initialize(create_mock_cluster_config(strategy="disaggregate"), medium_scheduling) + + assert manager.is_initialized is True + assert manager.get_placement_strategy() == "disaggregate" + assert set(manager.get_resource_pools()) == {"train_pool", "rollout_pool"} + assert manager.get_pool_for_component("train").name == "train_pool" + assert manager.get_pool_for_component("eval").name == "rollout_pool" + assert manager.get_storage_to_train_workers() == {0: list(range(8))} + assert manager.get_scheduling_strategy("train", 0) != "DEFAULT" + assert set(created) == {"train_pool_node_0", "rollout_pool_node_1"} + + +def test_global_resource_manager_initializes_manual_mode(monkeypatch, mock_cluster_2_nodes_8_gpus, basic_scheduling): + patch_cluster_sources(monkeypatch, mock_cluster_2_nodes_8_gpus) + patch_ray_runtime(monkeypatch) + + resource_pool = [ + { + "name": "train_pool", + "num_node": 1, + "num_gpus": 8, + "train": "0-7", + }, + { + "name": "rollout_pool", + "num_node": 1, + "num_gpus": 8, + "eval": "0-3", + "env": "4-7", + }, + ] + + manager = GlobalResourceManager.get_instance() + manager.initialize( + create_mock_cluster_config(mode="manual", resource_pool=resource_pool), + basic_scheduling, + ) + + assert manager.is_initialized is True + assert manager.get_placement_strategy() == "resource_pool" + assert set(manager.get_resource_pools()) == {"train_pool", "rollout_pool"} + assert manager.get_pool_for_component("train").name == "train_pool" + assert manager.get_pool_for_component("eval").name == "rollout_pool" + + +def test_global_resource_manager_saves_yaml_config(monkeypatch, tmp_path, mock_cluster_2_nodes_8_gpus, medium_scheduling): + patch_cluster_sources(monkeypatch, mock_cluster_2_nodes_8_gpus) + patch_ray_runtime(monkeypatch) + + manager = GlobalResourceManager.get_instance() + manager.initialize(create_mock_cluster_config(strategy="disaggregate"), medium_scheduling) + + saved_path = Path(manager.save_yaml_config(str(tmp_path))) + saved_config = yaml.safe_load(saved_path.read_text()) + + assert saved_path.exists() + assert saved_path.name == "resource_pool_auto.yaml" + assert len(saved_config) == 2 + assert saved_config[0]["name"] == "train_pool" + assert saved_config[1]["name"] == "rollout_pool" diff --git a/tests/unit/test_placement_strategies.py b/tests/unit/test_placement_strategies.py new file mode 100644 index 0000000..ac0e5f7 --- /dev/null +++ b/tests/unit/test_placement_strategies.py @@ -0,0 +1,321 @@ +from unittest.mock import MagicMock, patch + +import pytest + +from rlightning.utils.placement import ResourcePoolPlanner +from rlightning.utils.placement.placement_strategies import ( + ColocatedPlacementStrategy, + DefaultPlacementStrategy, + DisaggregatePlacementStrategy, + _pack_workers_on_gpu_units, +) +from rlightning.utils.placement.scheduling import ComponentScheduling, Scheduling + + +def create_mock_cluster_resources( + num_nodes: int = 2, + gpus_per_node: int = 8, + cpus_per_node: int = 64, +): + node_id_to_resources = {} + for i in range(num_nodes): + node_id = f"node_{i}" + node_id_to_resources[node_id] = { + "node_id": node_id, + "ip": f"192.168.1.{i + 1}", + "CPU": cpus_per_node, + "GPU": gpus_per_node, + } + return { + "node_id_to_resources": node_id_to_resources, + "total_cpus": num_nodes * cpus_per_node, + "total_gpus": num_nodes * gpus_per_node, + } + + +def create_mock_scheduling( + train_workers: int = 0, + train_gpus: float = 1.0, + eval_workers: int = 0, + eval_gpus: float = 1.0, + buffer_workers: int | str = 0, + buffer_gpus: float = 0.0, + env_workers: int = 0, + env_gpus: float = 1.0, +) -> ComponentScheduling: + return ComponentScheduling( + train_worker=Scheduling(worker_num=train_workers, num_cpus=1, num_gpus=train_gpus), + eval_worker=Scheduling(worker_num=eval_workers, num_cpus=1, num_gpus=eval_gpus), + buffer_worker=Scheduling(worker_num=buffer_workers, num_cpus=1, num_gpus=buffer_gpus), + env_worker=[Scheduling(worker_num=env_workers, num_cpus=1, num_gpus=env_gpus)], + ) + + +@pytest.fixture +def mock_cluster_2_nodes_8_gpus(): + return create_mock_cluster_resources(num_nodes=2, gpus_per_node=8) + + +@pytest.fixture +def mock_cluster_1_node_8_gpus(): + return create_mock_cluster_resources(num_nodes=1, gpus_per_node=8) + + +@pytest.fixture +def mock_cluster_3_nodes_8_gpus(): + return create_mock_cluster_resources(num_nodes=3, gpus_per_node=8) + + +@pytest.fixture +def basic_scheduling(): + return create_mock_scheduling( + train_workers=4, + train_gpus=1.0, + eval_workers=2, + eval_gpus=1.0, + buffer_workers=1, + env_workers=4, + env_gpus=0.5, + ) + + +@pytest.fixture +def large_train_scheduling(): + return create_mock_scheduling( + train_workers=16, + train_gpus=1.0, + eval_workers=4, + eval_gpus=1.0, + buffer_workers="auto", + env_workers=8, + env_gpus=0.5, + ) + + +def test_pack_workers_basic(): + worker_locations = [] + unit_cpu = [0, 0, 0, 0] + distribution = {} + + workers_placed = _pack_workers_on_gpu_units( + allocations=[(0, 3)], + node_id="node-0", + component_type="train", + node_component_distribution=distribution, + pg_key="test_pg", + capacity_gpus=4, + unit_cpu=unit_cpu, + worker_locations=worker_locations, + workers_placed=0, + workers_total=4, + gpu_req_list=[1.0, 1.0, 1.0, 1.0], + cpu_req=1, + ) + + assert workers_placed == 4 + assert len(worker_locations) == 4 + assert all(loc[0] == "test_pg" for loc in worker_locations) + assert distribution["node-0"]["train"] == {"count": 4, "ids": [0, 1, 2, 3]} + + +def test_pack_workers_fractional_gpu(): + worker_locations = [] + unit_cpu = [0, 0] + distribution = {} + + workers_placed = _pack_workers_on_gpu_units( + allocations=[(0, 1)], + node_id="node-0", + component_type="eval", + node_component_distribution=distribution, + pg_key="test_pg", + capacity_gpus=2, + unit_cpu=unit_cpu, + worker_locations=worker_locations, + workers_placed=0, + workers_total=4, + gpu_req_list=[0.5, 0.5, 0.5, 0.5], + cpu_req=1, + ) + + assert workers_placed == 4 + assert distribution["node-0"]["eval"]["count"] == 4 + + +def test_pack_workers_invalid_unit_index(): + worker_locations = [] + unit_cpu = [0, 0] + distribution = {} + + with pytest.raises(RuntimeError, match="Invalid GPU unit index"): + _pack_workers_on_gpu_units( + allocations=[(0, 5)], + node_id="node-0", + component_type="train", + node_component_distribution=distribution, + pg_key="test_pg", + capacity_gpus=2, + unit_cpu=unit_cpu, + worker_locations=worker_locations, + workers_placed=0, + workers_total=2, + gpu_req_list=[1.0, 1.0], + cpu_req=1, + ) + + +@patch("ray.util.placement_group") +@patch("ray.get") +@patch("rlightning.utils.placement.resource_pool.get_cluster_resources") +def test_disaggregate_create_placement_groups( + mock_pool_cluster, + mock_ray_get, + mock_placement_group, + basic_scheduling, + mock_cluster_2_nodes_8_gpus, +): + mock_pool_cluster.return_value = mock_cluster_2_nodes_8_gpus + mock_pg = MagicMock() + mock_pg.ready.return_value = True + mock_placement_group.return_value = mock_pg + mock_ray_get.return_value = None + + strategy = DisaggregatePlacementStrategy(basic_scheduling) + planner = ResourcePoolPlanner(scheduling=basic_scheduling) + pools = planner.plan_resource_pools(strategy="disaggregate") + + train_node_count = planner.get_component_node_count("train") + strategy.scheduling.adjust_buffer_worker_num(train_node_count) + pgs = strategy.create_placement_groups(resource_pools=pools) + + assert "train_pool_node_0" in pgs + assert strategy._storage_to_train_workers == {0: [0, 1, 2, 3]} + assert [loc[1] for loc in strategy._worker_locations["train"]] == [0, 1, 2, 3] + assert [loc[1] for loc in strategy._worker_locations["buffer"]] == [8] + assert [loc[1] for loc in strategy._worker_locations["eval"]] == [4, 5] + assert [loc[1] for loc in strategy._worker_locations["env"]] == [6, 6, 7, 7] + + +@patch("ray.util.placement_group") +@patch("ray.get") +@patch("rlightning.utils.placement.resource_pool.get_cluster_resources") +def test_colocated_create_placement_groups( + mock_pool_cluster, + mock_ray_get, + mock_placement_group, + basic_scheduling, + mock_cluster_1_node_8_gpus, +): + mock_pool_cluster.return_value = mock_cluster_1_node_8_gpus + mock_pg = MagicMock() + mock_pg.ready.return_value = True + mock_placement_group.return_value = mock_pg + mock_ray_get.return_value = None + + strategy = ColocatedPlacementStrategy(basic_scheduling) + planner = ResourcePoolPlanner(scheduling=basic_scheduling) + pools = planner.plan_resource_pools(strategy="colocate") + + train_node_count = planner.get_component_node_count("train") + strategy.scheduling.adjust_buffer_worker_num(train_node_count) + pgs = strategy.create_placement_groups(resource_pools=pools) + + assert strategy._storage_to_train_workers == {0: [0, 1, 2, 3]} + assert "global_pool_node_0" in pgs + assert [loc[1] for loc in strategy._worker_locations["train"]] == [0, 1, 2, 3] + assert [loc[1] for loc in strategy._worker_locations["buffer"]] == [4] + assert [loc[1] for loc in strategy._worker_locations["eval"]] == [0, 1] + assert [loc[1] for loc in strategy._worker_locations["env"]] == [2, 2, 3, 3] + + +@patch("rlightning.utils.placement.placement_strategies.get_cluster_resources") +def test_default_strategy_create_placement_groups_single_buffer( + mock_get_cluster, basic_scheduling, mock_cluster_2_nodes_8_gpus +): + mock_get_cluster.return_value = mock_cluster_2_nodes_8_gpus + + strategy = DefaultPlacementStrategy(basic_scheduling) + result = strategy.create_placement_groups() + + assert len(result) == 0 + + +@patch("rlightning.utils.placement.placement_strategies.get_cluster_resources") +def test_default_strategy_create_placement_groups_multiple_buffers(mock_get_cluster, mock_cluster_2_nodes_8_gpus): + mock_get_cluster.return_value = mock_cluster_2_nodes_8_gpus + scheduling = create_mock_scheduling(buffer_workers=2) + + strategy = DefaultPlacementStrategy(scheduling) + strategy.create_placement_groups() + + assert len(strategy.buffer_strategies) == 2 + + +@patch("rlightning.utils.placement.placement_strategies.get_cluster_resources") +def test_default_strategy_returns_default_scheduling(mock_get_cluster, basic_scheduling, mock_cluster_2_nodes_8_gpus): + mock_get_cluster.return_value = mock_cluster_2_nodes_8_gpus + + strategy = DefaultPlacementStrategy(basic_scheduling) + strategy.create_placement_groups() + + assert strategy.get_scheduling_strategy("train", 0) == "DEFAULT" + + +@patch("rlightning.utils.placement.placement_strategies.get_cluster_resources") +def test_default_strategy_keeps_empty_storage_mapping(mock_get_cluster, mock_cluster_2_nodes_8_gpus): + mock_get_cluster.return_value = mock_cluster_2_nodes_8_gpus + scheduling = create_mock_scheduling(train_workers=4, buffer_workers=2) + + strategy = DefaultPlacementStrategy(scheduling) + strategy.create_placement_groups() + + assert strategy.get_storage_to_train_workers() == {} + + +@patch("ray.util.placement_group") +@patch("ray.get") +@patch("rlightning.utils.placement.resource_pool.get_cluster_resources") +def test_disaggregate_train_workers_multi_node( + mock_pool_cluster, + mock_ray_get, + mock_placement_group, + mock_cluster_3_nodes_8_gpus, + large_train_scheduling, +): + mock_pool_cluster.return_value = mock_cluster_3_nodes_8_gpus + mock_pg = MagicMock() + mock_pg.ready.return_value = True + mock_placement_group.return_value = mock_pg + mock_ray_get.return_value = None + + scheduling = large_train_scheduling + strategy = DisaggregatePlacementStrategy(scheduling) + planner = ResourcePoolPlanner(scheduling=scheduling) + pools = planner.plan_resource_pools(strategy="disaggregate") + + train_node_count = planner.get_component_node_count("train") + scheduling.adjust_buffer_worker_num(train_node_count) + strategy.create_placement_groups(resource_pools=pools) + + assert train_node_count == 2 + train_locations = strategy._worker_locations["train"] + assert len(train_locations) == 16 + + train_pg_keys = {loc[0] for loc in train_locations} + assert len(train_pg_keys) == 2 + + node_0_locs = [loc for loc in train_locations if "node_0" in loc[0]] + node_1_locs = [loc for loc in train_locations if "node_1" in loc[0]] + assert len(node_0_locs) == 8 + assert len(node_1_locs) == 8 + + assert sorted(loc[1] for loc in node_0_locs) == list(range(8)) + assert sorted(loc[1] for loc in node_1_locs) == list(range(8)) + + buffer_locations = strategy._worker_locations["buffer"] + assert len(buffer_locations) == 2 + + mapping = strategy.get_storage_to_train_workers() + assert len(mapping) == 2 + assert mapping[0] == list(range(0, 8)) + assert mapping[1] == list(range(8, 16)) diff --git a/tests/unit/test_policy.py b/tests/unit/test_policy.py new file mode 100644 index 0000000..c6f1476 --- /dev/null +++ b/tests/unit/test_policy.py @@ -0,0 +1,23 @@ +import importlib +import inspect + +import pytest + +@pytest.mark.parametrize( + ("module_path", "class_name"), + [ + ("rlightning.policy.simple_ppo_policy.ppo_policy", "SimplePPOPolicy"), + ("rlightning.policy.rsl_rl_policy.rsl_rl_policy", "RSLRLPolicy"), + ("rlightning.policy.vla_policy.ppo_policy", "VLAPPOPolicy"), + ("rlightning.policy.supervised_policy", "SimpleSupervisedPolicy"), + ], +) +def test_policy_is_concrete(module_path: str, class_name: str): + try: + module = importlib.import_module(module_path) + except Exception as exc: + pytest.skip(f"Skipping {module_path} import: {exc}") + + policy_cls = getattr(module, class_name, None) + assert policy_cls is not None, f"Missing class {class_name} in {module_path}" + assert inspect.isabstract(policy_cls) is False diff --git a/tests/unit/test_policy_group.py b/tests/unit/test_policy_group.py new file mode 100644 index 0000000..f925733 --- /dev/null +++ b/tests/unit/test_policy_group.py @@ -0,0 +1,757 @@ +"""Unit tests for PolicyGroup and transfer-info helpers.""" + +from __future__ import annotations + +from collections import deque +from typing import Any, Dict, List + +import pytest +import torch + +from rlightning.policy.base_policy import PolicyRole +from rlightning.policy.policy_group import ( + PolicyGroup, + build_transfer_info_for_colocated, + build_transfer_info_for_disaggregated, +) +from rlightning.types import BatchedData, EnvRet, PolicyResponse +from rlightning.utils.config import PolicyConfig, TrainConfig, WeightBufferConfig +from rlightning.utils.distributed.group_initializer import ParallelMode + +@pytest.fixture(autouse=True) +def _reset_remote_flags(monkeypatch): + """Default unit tests to local mode unless a test opts into remote mode.""" + monkeypatch.setenv("RLIGHTNING_REMOTE_TRAIN", "0") + monkeypatch.setenv("RLIGHTNING_REMOTE_EVAL", "0") + + +class _DummyPolicy: + def __init__(self, name: str) -> None: + self.name = name + self.calls: Dict[str, List] = { + "init_eval": [], + "init_train": [], + "rollout": [], + "postprocess": [], + "update_dataset": [], + "train": [], + "print_timing_summary": [], + "notify_update_weights": 0, + "reset_training_state": [], + "send_weights": [], + "save_checkpoint": [], + "init_weight_buffer": [], + } + + def init_eval(self, eval_config=None, env_meta=None): + self.calls["init_eval"].append((eval_config, env_meta)) + return f"eval-{self.name}" + + def init_train(self, train_config=None, env_meta=None): + self.calls["init_train"].append((train_config, env_meta)) + return f"train-{self.name}" + + def _rollout(self, env_ret: EnvRet): + self.calls["rollout"].append(env_ret.env_id) + return PolicyResponse(env_id=env_ret.env_id, action=torch.tensor(1)) + + def _postprocess(self, *args: Any): + self.calls["postprocess"].append(args) + return (self.name, args) + + def update_dataset(self, data: Any): + self.calls["update_dataset"].append(data) + return f"update-{self.name}" + + def train(self, *args: Any, **kwargs: Any): + self.calls["train"].append((args, kwargs)) + return {"policy": self.name} + + def print_timing_summary(self, reset: bool = False): + self.calls["print_timing_summary"].append(reset) + + def notify_update_weights(self): + self.calls["notify_update_weights"] += 1 + + def reset_training_state(self, train_config, env_meta=None, seed=None): + self.calls["reset_training_state"].append((train_config, env_meta, seed)) + + def send_weights(self, receivers, shared_weight_buffer=None): + self.calls["send_weights"].append((receivers, shared_weight_buffer)) + + def save_checkpoint(self, save_dir): + self.calls["save_checkpoint"].append(save_dir) + + def init_weight_buffer(self, shared_weight_buffer=None): + self.calls["init_weight_buffer"].append(shared_weight_buffer) + + +class _FakeObjectRef: + def __init__(self, value=None, hex_value: str | None = None) -> None: + self.value = value + self._hex = hex_value or hex(id(self)) + + def hex(self) -> str: + return self._hex + + +class _RemoteMethod: + def __init__(self, fn, wrap_ref: bool = False): + self.fn = fn + self.wrap_ref = wrap_ref + self.calls: List = [] + + def __call__(self, *args, **kwargs): + return self.fn(*args, **kwargs) + + def remote(self, *args, **kwargs): + self.calls.append((args, kwargs)) + result = self.fn(*args, **kwargs) + if self.wrap_ref: + return _FakeObjectRef(result) + return result + + +class _AsyncPolicy: + def __init__(self, name: str, num_requests: int) -> None: + self.name = name + self._num_requests = num_requests + self.rollout_calls: List[str] = [] + self.get_num_requests = _RemoteMethod(self._get_num_requests) + self._rollout_async = _RemoteMethod(self._rollout) + + def _get_num_requests(self): + return self._num_requests + + def _rollout(self, env_ret: EnvRet): + self.rollout_calls.append(env_ret.env_id) + return f"resp-{self.name}-{env_ret.env_id}" + + +def _make_policy_cfg( + *, + rollout_mode: str = "sync", + buffer_strategy: str = "Double", + buffer_type: str = "WeightBuffer", + is_colocated: bool = False, +) -> PolicyConfig: + return PolicyConfig( + type="DummyPolicy", + rollout_mode=rollout_mode, + is_colocated=is_colocated, + weight_buffer=WeightBufferConfig(type=buffer_type, buffer_strategy=buffer_strategy), + ) + + +def test_policy_group_init_calls_eval_then_train(): + """init should call init_eval and init_train for all policies.""" + train_policy = _DummyPolicy("train") + eval_policy = _DummyPolicy("eval") + group = PolicyGroup([train_policy], [eval_policy], _make_policy_cfg()) + train_cfg = TrainConfig(max_epochs=1) + + group.init(train_cfg, env_meta="meta", eval_config="eval-cfg") + + assert eval_policy.calls["init_eval"] == [("eval-cfg", "meta")] + assert train_policy.calls["init_train"] == [(train_cfg, "meta")] + + +def test_policy_group_rollout_batch_sync_uses_idle_deque(): + """rollout_batch should dispatch sync rollouts via idle deque.""" + eval_policy = _DummyPolicy("eval-1") + group = PolicyGroup([], [eval_policy], _make_policy_cfg(rollout_mode="sync")) + group._idle_deque = deque() + + env_rets = { + "env-1": EnvRet(env_id="env-1", observation=0), + "env-2": EnvRet(env_id="env-2", observation=1), + } + batched = BatchedData.from_dict(env_rets) + + result = group.rollout_batch(batched) + + assert result.ids() == ("env-1", "env-2") + assert [resp.env_id for resp in result.values()] == ["env-1", "env-2"] + assert eval_policy.calls["rollout"] == ["env-1", "env-2"] + + +def test_policy_group_rollout_batch_async_uses_router(monkeypatch): + """rollout_batch should route async requests based on router assignments.""" + policies = [_AsyncPolicy("p0", num_requests=2), _AsyncPolicy("p1", num_requests=0)] + group = PolicyGroup([], policies, _make_policy_cfg(rollout_mode="async")) + + assign_calls: Dict[str, List[int]] = {} + + def fake_assign(loads, n): + assign_calls["loads"] = loads + assign_calls["n"] = n + return [1, 0] + + group.router.assign = fake_assign + monkeypatch.setattr("rlightning.policy.policy_group.ray.get", lambda x: x) + + env_rets = { + "env-1": EnvRet(env_id="env-1", observation=0), + "env-2": EnvRet(env_id="env-2", observation=1), + } + batched = BatchedData.from_dict(env_rets) + + result = group.rollout_batch(batched) + + assert assign_calls["loads"] == [2, 0] + assert assign_calls["n"] == 2 + assert policies[1].rollout_calls == ["env-1"] + assert policies[0].rollout_calls == ["env-2"] + assert result.values() == ("resp-p1-env-1", "resp-p0-env-2") + + +def test_policy_group_postprocess_sync_dispatches(): + """postprocess should dispatch to eval policies and preserve ids.""" + eval_policy = _DummyPolicy("eval-1") + group = PolicyGroup([], [eval_policy], _make_policy_cfg()) + group._idle_deque = deque() + + env_rets = { + "env-1": EnvRet(env_id="env-1", observation=0), + "env-2": EnvRet(env_id="env-2", observation=1), + } + policy_resps = { + "env-1": PolicyResponse(env_id="env-1", action=1), + "env-2": PolicyResponse(env_id="env-2", action=2), + } + + batched_env = BatchedData.from_dict(env_rets) + batched_resp = BatchedData.from_dict(policy_resps) + + result = group.postprocess(batched_env, batched_resp) + + assert result.ids() == ("env-1", "env-2") + assert len(result.values()) == 2 + assert eval_policy.calls["postprocess"] + + +def test_policy_group_postprocess_env_only(): + """postprocess should work with only batched_env_ret.""" + eval_policy = _DummyPolicy("eval-1") + group = PolicyGroup([], [eval_policy], _make_policy_cfg()) + group._idle_deque = deque([eval_policy]) + + env_rets = {"env-1": EnvRet(env_id="env-1", observation=0)} + batched_env = BatchedData.from_dict(env_rets) + + result = group.postprocess(batched_env, None) + + assert result.ids() == ("env-1",) + name, args = result.values()[0] + assert name == "eval-1" + assert len(args) == 1 + assert isinstance(args[0], EnvRet) + + +def test_policy_group_postprocess_policy_only(): + """postprocess should work with only batched_policy_resp.""" + eval_policy = _DummyPolicy("eval-1") + group = PolicyGroup([], [eval_policy], _make_policy_cfg()) + group._idle_deque = deque([eval_policy]) + + policy_resps = {"env-1": PolicyResponse(env_id="env-1", action=1)} + batched_resp = BatchedData.from_dict(policy_resps) + + result = group.postprocess(None, batched_resp) + + assert result.ids() == ("env-1",) + name, args = result.values()[0] + assert name == "eval-1" + assert len(args) == 1 + assert isinstance(args[0], PolicyResponse) + + +def test_policy_group_update_dataset_and_train(): + """update_dataset and train should route to all train policies.""" + train_policies = [_DummyPolicy("t1"), _DummyPolicy("t2")] + group = PolicyGroup(train_policies, [], _make_policy_cfg()) + + updates = group.update_dataset(["d1", "d2"]) + assert updates == ["update-t1", "update-t2"] + assert train_policies[0].calls["update_dataset"] == ["d1"] + + train_info = group.train(123) + assert train_info == {"policy": "t1"} + + +def test_policy_group_update_dataset_mismatched_length_raises(): + """update_dataset should assert when length mismatches train policies.""" + train_policies = [_DummyPolicy("t1")] + group = PolicyGroup(train_policies, [], _make_policy_cfg()) + + with pytest.raises(AssertionError): + group.update_dataset([]) + + +def test_policy_group_reset_training_state_calls_policies(): + """reset_training_state should call all train policies.""" + train_policies = [_DummyPolicy("t1"), _DummyPolicy("t2")] + group = PolicyGroup(train_policies, [], _make_policy_cfg()) + train_cfg = TrainConfig(max_epochs=1) + + group.reset_training_state(train_cfg, env_meta="meta", seed=123) + + assert train_policies[0].calls["reset_training_state"] == [(train_cfg, "meta", 123)] + assert train_policies[1].calls["reset_training_state"] == [(train_cfg, "meta", 123)] + + +def test_policy_group_reset_training_state_remote(monkeypatch): + """reset_training_state should ray.get refs and call dist_barrier for remote refs.""" + train_policy = _DummyPolicy("t1") + train_policy.dist_barrier = _RemoteMethod(lambda *args: "barrier", wrap_ref=True) + group = PolicyGroup([train_policy], [], _make_policy_cfg()) + train_cfg = TrainConfig(max_epochs=1) + + def fake_submit(method, *args, _block=False, **kwargs): + result = method(*args, **kwargs) + if _block: + return result + return _FakeObjectRef(result) + + def fake_ray_get(obj): + if isinstance(obj, list): + return [fake_ray_get(o) for o in obj] + if isinstance(obj, _FakeObjectRef): + return obj.value + return obj + + monkeypatch.setattr("rlightning.policy.policy_group.ray.ObjectRef", _FakeObjectRef, raising=False) + monkeypatch.setattr("rlightning.policy.policy_group.ray.get", fake_ray_get) + + group._eval_task_submitter.submit = fake_submit + group._train_task_submitter.submit = fake_submit + group.reset_training_state(train_cfg, env_meta="meta", seed=321) + + assert train_policy.calls["reset_training_state"] == [(train_cfg, "meta", 321)] + assert train_policy.dist_barrier.calls[0][0][0] == ParallelMode.TRAIN_DATA_PARALLEL + + +def test_policy_group_print_timing_summary_calls_policies(): + """print_timing_summary should call all policies with reset flag.""" + train_policy = _DummyPolicy("train") + eval_policy = _DummyPolicy("eval") + group = PolicyGroup([train_policy], [eval_policy], _make_policy_cfg()) + + group.print_timing_summary(reset=True) + + assert train_policy.calls["print_timing_summary"] == [True] + assert eval_policy.calls["print_timing_summary"] == [True] + + +def test_policy_group_init_eval_and_train_calls(): + """init_eval and init_train should call policy init methods.""" + train_policy = _DummyPolicy("train") + eval_policy = _DummyPolicy("eval") + group = PolicyGroup([train_policy], [eval_policy], _make_policy_cfg()) + + group.init_eval(eval_config="cfg", env_meta="meta") + group.init_train(train_config=TrainConfig(max_epochs=1), env_meta="meta") + + assert train_policy.calls["init_train"] + assert eval_policy.calls["init_eval"] + + +def test_policy_group_init_comm_group_empty_is_noop(): + """init_comm_group should return early when no policies exist.""" + group = PolicyGroup([], [], _make_policy_cfg()) + group.init_comm_group(backend="gloo") + + +def test_policy_group_init_comm_group_remote_with_transfer(monkeypatch): + """init_comm_group should wire up remote groups and barriers.""" + + class _RemotePolicy: + def __init__( + self, + name: str, + node_id: str, + gpu_id: int, + rank: int, + addr_port: tuple[str, int], + ) -> None: + self.name = name + self._get_node_id = _RemoteMethod(lambda: node_id) + self._get_gpu_ids = _RemoteMethod(lambda: [gpu_id]) + self._get_addr_and_port = _RemoteMethod(lambda: addr_port, wrap_ref=True) + self.init_distributed_env = _RemoteMethod(lambda **kwargs: kwargs, wrap_ref=True) + self.get_rank = _RemoteMethod(lambda: rank, wrap_ref=True) + self.init_single_comm_group = _RemoteMethod( + lambda ranks, mode, backend: (ranks, mode, backend), wrap_ref=True + ) + self.dist_barrier = _RemoteMethod(lambda *args: "barrier", wrap_ref=True) + + train_policy = _RemotePolicy("train", "node-a", 0, 0, ("1.2.3.4", 1234)) + eval_policy = _RemotePolicy("eval", "node-b", 1, 1, ("9.9.9.9", 9999)) + + group = PolicyGroup([train_policy], [eval_policy], _make_policy_cfg()) + group.weight_buffer_strategy = "Double" + group.transfer_info = { + "node-a": [ + { + "sender": train_policy, + "receiver": [eval_policy], + "intra_node": False, + } + ] + } + + def fake_ray_get(obj): + if isinstance(obj, list): + return [fake_ray_get(o) for o in obj] + if isinstance(obj, _FakeObjectRef): + return obj.value + return obj + + monkeypatch.setattr("rlightning.policy.policy_group.ray.get", fake_ray_get) + group.init_comm_group(backend="gloo") + + assert len(train_policy.init_distributed_env.calls) == 1 + assert len(eval_policy.init_distributed_env.calls) == 1 + + _, train_kwargs = train_policy.init_distributed_env.calls[0] + _, eval_kwargs = eval_policy.init_distributed_env.calls[0] + + assert train_kwargs["world_size"] == 2 + assert eval_kwargs["world_size"] == 2 + assert train_kwargs["master_addr"] == "1.2.3.4" + assert train_kwargs["master_port"] == 1234 + assert train_kwargs["rank"] == 0 + assert eval_kwargs["rank"] == 1 + assert train_kwargs["local_rank"] == 0 + assert eval_kwargs["local_rank"] == 0 + assert train_kwargs["local_world_size"] == 1 + assert eval_kwargs["local_world_size"] == 1 + + for policy in (train_policy, eval_policy): + modes = [args[1] for args, _ in policy.init_single_comm_group.calls] + assert ParallelMode.INTRA_NODE in modes + assert ParallelMode.WEIGHT_TRANSFER in modes + assert ParallelMode.TRAIN_DATA_PARALLEL in modes + + transfer_calls = [ + args for args, _ in train_policy.init_single_comm_group.calls if args[1] == ParallelMode.WEIGHT_TRANSFER + ] + assert transfer_calls[0][0] == [0, 1] + + assert len(train_policy.dist_barrier.calls) == 3 + assert len(eval_policy.dist_barrier.calls) == 3 + + +def test_policy_group_push_pop_updates_lists(): + """push/pop should update role lists.""" + group = PolicyGroup([], [], _make_policy_cfg()) + train_policy = _DummyPolicy("train") + eval_policy = _DummyPolicy("eval") + + group.push(train_policy, PolicyRole.TRAIN) + group.push(eval_policy, PolicyRole.EVAL) + + assert group.train_list[-1] is train_policy + assert group.eval_list[-1] is eval_policy + + popped = group.pop(PolicyRole.EVAL) + assert popped is eval_policy + assert group.eval_list == [] + + +def test_policy_group_init_placement_info_local(monkeypatch): + """init_placement_info should assign local node for local policies.""" + train_policy = _DummyPolicy("train") + eval_policy = _DummyPolicy("eval") + group = PolicyGroup([train_policy], [eval_policy], _make_policy_cfg(buffer_strategy="Double")) + + monkeypatch.setenv("RLIGHTNING_REMOTE_TRAIN", "0") + monkeypatch.setenv("RLIGHTNING_REMOTE_EVAL", "0") + + group.init_placement_info() + + assert "local" in group.placement_info + assert train_policy in group.placement_info["local"]["train_policies"] + assert eval_policy in group.placement_info["local"]["eval_policies"] + assert group.transfer_info["local"][0]["sender"] is train_policy + assert group.transfer_info["local"][0]["receiver"] == [eval_policy] + assert group.transfer_info["local"][0]["intra_node"] is True + + +def test_policy_group_init_placement_info_remote(monkeypatch): + """init_placement_info should query node ids via ray for remote policies.""" + + class _RemotePolicy: + def __init__(self, name: str, node_id: str) -> None: + self.name = name + self._get_node_id = _RemoteMethod(lambda: node_id) + + train_policy = _RemotePolicy("train", "node-train") + eval_policy = _RemotePolicy("eval", "node-eval") + group = PolicyGroup([train_policy], [eval_policy], _make_policy_cfg(buffer_strategy="Double")) + + monkeypatch.setenv("RLIGHTNING_REMOTE_TRAIN", "1") + monkeypatch.setenv("RLIGHTNING_REMOTE_EVAL", "1") + + def fake_ray_get(obj): + if isinstance(obj, list): + return [fake_ray_get(o) for o in obj] + if isinstance(obj, _FakeObjectRef): + return obj.value + return obj + + monkeypatch.setattr("rlightning.policy.policy_group.ray.get", fake_ray_get) + group.init_placement_info() + + assert "node-train" in group.placement_info + assert "node-eval" in group.placement_info + assert train_policy in group.placement_info["node-train"]["train_policies"] + assert eval_policy in group.placement_info["node-eval"]["eval_policies"] + assert group.transfer_info["node-train"][0]["receiver"] == [eval_policy] + assert group.transfer_info["node-train"][0]["intra_node"] is False + + +def test_policy_group_init_placement_info_colocated(): + """init_placement_info should use colocated helper when configured.""" + train_policy = _DummyPolicy("train") + eval_policy = _DummyPolicy("eval") + group = PolicyGroup( + [train_policy], + [eval_policy], + _make_policy_cfg(buffer_strategy="None"), + ) + + group.init_placement_info(is_colocated=True) + + assert group.transfer_info == {"all": [{"sender": train_policy, "receiver": eval_policy}]} + assert group.placement_info == {} + + +def test_policy_group_init_weight_buffer_double(): + """init_weight_buffer should call init on eval policies for Double strategy.""" + eval_policy = _DummyPolicy("eval") + group = PolicyGroup([], [eval_policy], _make_policy_cfg(buffer_strategy="Double")) + + group.init_weight_buffer() + + assert eval_policy.calls["init_weight_buffer"] == [None] + + +def test_policy_group_init_weight_buffer_none_noop(): + """init_weight_buffer should skip when strategy is None.""" + eval_policy = _DummyPolicy("eval") + group = PolicyGroup([], [eval_policy], _make_policy_cfg(buffer_strategy="None")) + + group.init_weight_buffer() + + assert eval_policy.calls["init_weight_buffer"] == [] + + +def test_policy_group_init_weight_buffer_invalid_strategy_raises(): + """init_weight_buffer should reject unsupported strategies.""" + eval_policy = _DummyPolicy("eval") + group = PolicyGroup([], [eval_policy], _make_policy_cfg(buffer_strategy="Double")) + group.config.weight_buffer.buffer_strategy = "BadStrategy" + + with pytest.raises(ValueError): + group.init_weight_buffer() + + +def test_policy_group_init_weight_buffer_shared(monkeypatch): + """Shared strategy should create shared buffer and reuse across eval policies.""" + eval_policy = _DummyPolicy("eval") + group = PolicyGroup( + [], + [eval_policy], + _make_policy_cfg(buffer_strategy="Shared", buffer_type="CPUWeightBuffer"), + ) + group.placement_info = {"local": {"train_policies": [], "eval_policies": [eval_policy]}} + + monkeypatch.setattr( + "rlightning.policy.policy_group.build_weight_buffer", + lambda *args, **kwargs: "shared-buffer", + ) + + group.init_weight_buffer() + + assert eval_policy.calls["init_weight_buffer"] == ["shared-buffer"] + assert group.shared_weight_buffer_map["local"]["shared_weight_buffer"] == "shared-buffer" + + +def test_build_transfer_info_for_disaggregated_various_strategies(): + """build_transfer_info_for_disaggregated should route train/eval nodes by strategy.""" + train = _DummyPolicy("train") + eval_a = _DummyPolicy("eval-a") + eval_b = _DummyPolicy("eval-b") + + placement = { + "node-train": {"train_policies": [train], "eval_policies": []}, + "node-eval": {"train_policies": [], "eval_policies": [eval_a, eval_b]}, + } + + transfer_info = build_transfer_info_for_disaggregated(placement, "Double") + assert transfer_info["node-train"][0]["receiver"] == [eval_a, eval_b] + assert transfer_info["node-train"][0]["intra_node"] is False + + transfer_info = build_transfer_info_for_disaggregated(placement, "Shared") + assert transfer_info["node-train"][0]["receiver"] == [eval_a] + + transfer_info = build_transfer_info_for_disaggregated(placement, "None") + assert transfer_info["node-train"][0]["receiver"] == [eval_a, eval_b] + + with pytest.raises(ValueError): + build_transfer_info_for_disaggregated( + {"node-eval": {"train_policies": [], "eval_policies": [eval_a]}}, + "Double", + ) + + +def test_build_transfer_info_for_colocated(): + """build_transfer_info_for_colocated should map train/eval policy pairs by index.""" + train_policies = [_DummyPolicy("train-0"), _DummyPolicy("train-1")] + eval_policies = [_DummyPolicy("eval-0"), _DummyPolicy("eval-1")] + + transfer_info = build_transfer_info_for_colocated(train_policies, eval_policies) + + assert transfer_info["all"][0]["sender"] is train_policies[0] + assert transfer_info["all"][0]["receiver"] is eval_policies[0] + assert transfer_info["all"][1]["sender"] is train_policies[1] + assert transfer_info["all"][1]["receiver"] is eval_policies[1] + + +def test_policy_group_send_weights_double(): + """send_weights should call train send_weights without shared buffer.""" + train_policy = _DummyPolicy("train") + eval_policy = _DummyPolicy("eval") + group = PolicyGroup([train_policy], [eval_policy], _make_policy_cfg()) + group.weight_buffer_strategy = "Double" + group.transfer_info = { + "local": [ + { + "sender": train_policy, + "receiver": [eval_policy], + "intra_node": False, + } + ] + } + + group.send_weights() + + assert train_policy.calls["send_weights"] == [([eval_policy], None)] + assert train_policy.calls["save_checkpoint"] == [] + + +def test_policy_group_send_weights_remote_shared(monkeypatch): + """send_weights should use remote calls when RLIGHTNING_REMOTE_TRAIN is set.""" + + class _RemoteTrainPolicy: + def __init__(self) -> None: + self.send_weights = _RemoteMethod( + lambda receivers, shared_weight_buffer=None: "sent", + wrap_ref=True, + ) + self.save_weights = _RemoteMethod(lambda save_dir, epoch: "saved", wrap_ref=True) + + train_policy = _RemoteTrainPolicy() + eval_policy = _DummyPolicy("eval") + + group = PolicyGroup([train_policy], [eval_policy], _make_policy_cfg()) + group.weight_buffer_strategy = "Shared" + group.transfer_info = { + "local": [ + { + "sender": train_policy, + "receiver": [eval_policy], + "intra_node": True, + } + ] + } + group.shared_weight_buffer_map = {"local": {"shared_weight_buffer": "buf"}} + + monkeypatch.setenv("RLIGHTNING_REMOTE_TRAIN", "1") + + def fake_ray_get(obj): + if isinstance(obj, list): + return [fake_ray_get(o) for o in obj] + if isinstance(obj, _FakeObjectRef): + return obj.value + return obj + + monkeypatch.setattr("rlightning.policy.policy_group.ray.get", fake_ray_get) + group.send_weights() + + assert train_policy.send_weights.calls[0][0] == ([eval_policy], "buf") + + +def test_policy_group_send_weights_shared(): + """send_weights should use shared buffer when configured for local mode.""" + train_policy = _DummyPolicy("train") + eval_policy = _DummyPolicy("eval") + + group = PolicyGroup([train_policy], [eval_policy], _make_policy_cfg()) + group.weight_buffer_strategy = "Shared" + group.transfer_info = { + "local": [ + { + "sender": train_policy, + "receiver": [eval_policy], + "intra_node": True, + } + ] + } + group.shared_weight_buffer_map = {"local": {"shared_weight_buffer": "buf"}} + + group.send_weights() + + assert train_policy.calls["send_weights"][0] == ([eval_policy], "buf") + + +def test_policy_group_save_checkpoint_local(tmp_path): + """save_checkpoint should call local save_checkpoint directly.""" + train_policy = _DummyPolicy("train") + group = PolicyGroup([train_policy], [], _make_policy_cfg()) + + group.save_checkpoint(str(tmp_path)) + + assert train_policy.calls["save_checkpoint"] == [str(tmp_path)] + + +def test_policy_group_save_checkpoint_remote(monkeypatch, tmp_path): + """save_checkpoint should call remote save_checkpoint when remote train is enabled.""" + + class _RemoteTrainPolicy: + def __init__(self) -> None: + self.save_checkpoint = _RemoteMethod(lambda path: "saved", wrap_ref=True) + + train_policy = _RemoteTrainPolicy() + group = PolicyGroup([train_policy], [], _make_policy_cfg()) + monkeypatch.setenv("RLIGHTNING_REMOTE_TRAIN", "1") + + def fake_submit(method, *args, _block: bool = False, **kwargs): + # Simulate TaskSubmitter behavior for remote actors by always calling .remote + return method.remote(*args, **kwargs) + + def fake_ray_get(obj): + if isinstance(obj, _FakeObjectRef): + return obj.value + return obj + + # Route PolicyGroup.save_checkpoint through our fake submitter and fake ray.get + group._eval_task_submitter.submit = fake_submit + group._train_task_submitter.submit = fake_submit + monkeypatch.setattr("rlightning.policy.policy_group.ray.get", fake_ray_get) + group.save_checkpoint(str(tmp_path)) + + # _RemoteMethod records calls as (args, kwargs); first arg of first call is the path + assert train_policy.save_checkpoint.calls[0][0][0] == str(tmp_path) + + +def test_policy_group_notify_update_weights(): + """notify_update_weights should signal all eval policies.""" + eval_policies = [_DummyPolicy("e1"), _DummyPolicy("e2")] + group = PolicyGroup([], eval_policies, _make_policy_cfg()) + + group.notify_update_weights() + + assert eval_policies[0].calls["notify_update_weights"] == 1 + assert eval_policies[1].calls["notify_update_weights"] == 1 diff --git a/tests/unit/test_policy_param_sync.py b/tests/unit/test_policy_param_sync.py new file mode 100644 index 0000000..263b6e4 --- /dev/null +++ b/tests/unit/test_policy_param_sync.py @@ -0,0 +1,66 @@ +from types import SimpleNamespace +import sys +import types + +import torch +import torch.nn as nn + +from rlightning.policy.base_policy import PolicyRole +from rlightning.policy.rsl_rl_policy import RSLRLPolicy, RSLRLVecEnvMeta +from rlightning.utils.config import Config, PolicyConfig, WeightBufferConfig + + +def _make_env_ret(observation): + return SimpleNamespace(env_id="env-0", observation=observation, info={}) + + +def _sync_and_compare(train_policy, eval_policy, env_ret): + train_action = train_policy.rollout_step(env_ret).action + eval_action = eval_policy.rollout_step(env_ret).action + assert not torch.allclose(train_action, eval_action) + + eval_policy.load_state_dict(train_policy.get_trainable_parameters()) + synced_action = eval_policy.rollout_step(env_ret).action + torch.testing.assert_close(synced_action, train_action) + + +def test_param_sync_rsl_rl_policy(monkeypatch): + class FakeAgent: + def __init__(self, env_meta, **kwargs): + self.actor = nn.Linear(3, 2, bias=False) + self.initialized = True + + def to(self, device): + self.actor.to(device) + return self + + def draw_actions(self, obs, info): + return self.actor(obs), {"data": torch.zeros(obs.shape[0], device=obs.device)} + + def draw_random_actions(self, obs, info): + return torch.zeros((obs.shape[0], 2), device=obs.device), { + "data": torch.zeros(obs.shape[0], device=obs.device) + } + + fake_algorithms = types.SimpleNamespace(PPO=FakeAgent) + fake_rsl_rl = types.SimpleNamespace(algorithms=fake_algorithms) + monkeypatch.setitem(sys.modules, "rsl_rl", fake_rsl_rl) + + policy_cfg = PolicyConfig( + type="RSLRLPolicy", + weight_buffer=WeightBufferConfig(type="WeightBuffer", buffer_strategy="Double"), + policy_kwargs=Config.from_dict({"algorithm": "PPO"}), + ) + train_policy = RSLRLPolicy(policy_cfg, PolicyRole.TRAIN) + eval_policy = RSLRLPolicy(policy_cfg, PolicyRole.EVAL) + + env_meta = RSLRLVecEnvMeta(num_envs=1, num_actions=2, get_observations=None) + train_policy.construct_network(env_meta=env_meta) + eval_policy.construct_network(env_meta=env_meta) + + with torch.no_grad(): + train_policy.algo.actor.weight.fill_(1.0) + eval_policy.algo.actor.weight.zero_() + + observation = torch.ones(1, 3, device=train_policy.device) + _sync_and_compare(train_policy, eval_policy, _make_env_ret(observation)) diff --git a/tests/unit/test_preprocessors.py b/tests/unit/test_preprocessors.py new file mode 100644 index 0000000..31d4553 --- /dev/null +++ b/tests/unit/test_preprocessors.py @@ -0,0 +1,127 @@ +"""Unit tests for buffer preprocessors.""" + +from __future__ import annotations + +from typing import Dict + +import gymnasium as gym +import numpy as np +import pytest +import torch + +from rlightning.buffer.utils.preprocessors import ( + BoxFlattenPreprocessor, + DiscretePreprocessor, + NonPreprocessor, + default_obs_preprocessor, + default_reward_preprocessor, + get_preprocessor_cls, +) + +def test_preprocessor_call_rejects_invalid_type(): + """Preprocessor __call__ should assert on unsupported input types. + But origin function uses assert, which will be ignored when python + is run with -O flag. So maybe the origin function should raise + TypeError instead. TODO: @yangzhenyu @qiujiawei + """ + space = gym.spaces.Box(low=-1.0, high=1.0, shape=(2,)) + pre = NonPreprocessor(space) + + with pytest.raises(AssertionError): + _ = pre(123) + + +def test_non_preprocessor_returns_input_and_shape(): + """NonPreprocessor should return inputs unchanged and expose shape.""" + space = gym.spaces.Box(low=-1.0, high=1.0, shape=(2, 3)) + pre = NonPreprocessor(space) + data = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) + # simple origin function, no need to test multiple types. + + assert pre.shape == (2, 3) + assert pre.transform(data) is data + assert pre.batch_transform(data) is data + + +def test_box_flatten_preprocessor_transform_and_batch(): + """BoxFlattenPreprocessor should flatten single and batched observations.""" + space = gym.spaces.Box(low=-1.0, high=1.0, shape=(2, 2)) + pre = BoxFlattenPreprocessor(space) + + assert pre.shape == (4,) + + torch_obs = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + flat_torch = pre.transform(torch_obs) + assert flat_torch.shape == (4,) + assert flat_torch.tolist() == [1.0, 2.0, 3.0, 4.0] + + np_obs = np.array([[5.0, 6.0], [7.0, 8.0]]) + flat_np = pre.transform(np_obs) + assert flat_np.shape == (4,) + assert flat_np.tolist() == [5.0, 6.0, 7.0, 8.0] + + batched = torch.tensor( + [ + [[1.0, 2.0], [3.0, 4.0]], + [[5.0, 6.0], [7.0, 8.0]], + ] + ) + flat_batch = pre.batch_transform(batched) + assert flat_batch.shape == (2, 4) + assert flat_batch.tolist() == [[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]] + + tuple_batch = (np.array([[1.0, 2.0], [3.0, 4.0]]), np.array([[9.0, 10.0], [11.0, 12.0]])) + tuple_out = pre.batch_transform(tuple_batch) + assert isinstance(tuple_out, tuple) + assert tuple_out[0].tolist() == [1.0, 2.0, 3.0, 4.0] + assert tuple_out[1].tolist() == [9.0, 10.0, 11.0, 12.0] + + +def test_discrete_preprocessor_one_hot_and_batch_tuple(): + """DiscretePreprocessor should one-hot encode inputs and support tuple batches.""" + space = gym.spaces.Discrete(3) + pre = DiscretePreprocessor(space) + + assert pre.shape == (3,) + + np_obs = np.array(1) + one_hot_np = pre.transform(np_obs) + assert isinstance(one_hot_np, np.ndarray) + assert one_hot_np.tolist() == [0.0, 1.0, 0.0] + + torch_obs = torch.tensor(2) + one_hot_torch = pre.transform(torch_obs) + assert isinstance(one_hot_torch, torch.Tensor) + assert one_hot_torch.tolist() == [0.0, 0.0, 1.0] + + tuple_batch = (np.array(0), np.array(2)) + tuple_out = pre.batch_transform(tuple_batch) + assert isinstance(tuple_out, tuple) + assert tuple_out[0].tolist() == [1.0, 0.0, 0.0] + assert tuple_out[1].tolist() == [0.0, 0.0, 1.0] + + tensor_batch = torch.tensor([0, 2]) + tensor_batch_out = pre.batch_transform(tensor_batch) + assert isinstance(tensor_batch_out, torch.Tensor) + assert tensor_batch_out.shape == (2, 3) + assert tensor_batch_out.tolist() == [[1.0, 0.0, 0.0], [0.0, 0.0, 1.0]] + + +def test_default_preprocessors_return_input_sequences(): + """default_obs_preprocessor and default_reward_preprocessor should pass through.""" + obs_seq: Dict[str, int] = {"a": 1} + rew_seq = [1.0, 2.0] + + assert default_obs_preprocessor(obs_seq) is obs_seq + assert default_reward_preprocessor(rew_seq) is rew_seq + + +def test_get_preprocessor_cls_selects_by_space_type(): + """get_preprocessor_cls should select the right preprocessor class.""" + box_space = gym.spaces.Box(low=-1.0, high=1.0, shape=(2,)) + discrete_space = gym.spaces.Discrete(5) + other_space = gym.spaces.MultiDiscrete([2, 3]) + + assert get_preprocessor_cls(box_space) is BoxFlattenPreprocessor + assert get_preprocessor_cls(discrete_space) is DiscretePreprocessor + assert get_preprocessor_cls(other_space) is NonPreprocessor diff --git a/tests/unit/test_profiler.py b/tests/unit/test_profiler.py new file mode 100644 index 0000000..69b1cdd --- /dev/null +++ b/tests/unit/test_profiler.py @@ -0,0 +1,217 @@ +""" +Tests for the profiler module. + +This module contains unit tests for the profiler utilities including timing +context managers, decorators, and latency computation functions. +""" + +import time +from unittest.mock import patch + +import pytest + +from rlightning.utils.profiler.profiler import record_timing, timer, timer_wrap + + +@pytest.fixture(autouse=True) +def _enable_profiler(monkeypatch): + monkeypatch.setenv("RLIGHTNING_DEBUG", "1") + + +class TestSimpleTimer: + """Test cases for the timer context manager.""" + + def test_timer_basic_usage(self): + """Test basic functionality of timer context manager.""" + timing_data = {} + + with timer("test_function", timing_data, level="info"): + time.sleep(0.01) # Sleep for 10ms + + # Verify timing data was recorded + assert "test_function" in timing_data + assert timing_data["test_function"]["count"] == 1 + assert timing_data["test_function"]["total"] > 0 + assert timing_data["test_function"]["avg"] > 0 + + def test_timer_multiple_calls(self): + """Test timer with multiple calls to same function.""" + timing_data = {} + + # First call + with timer("test_function", timing_data, level="info"): + time.sleep(0.01) + + # Second call + with timer("test_function", timing_data, level="info"): + time.sleep(0.02) + + # Verify cumulative statistics + assert timing_data["test_function"]["count"] == 2 + assert timing_data["test_function"]["total"] > 0.03 + assert timing_data["test_function"]["avg"] > 0.015 + + def test_timer_different_functions(self): + """Test timer with different function names.""" + timing_data = {} + + with timer("func1", timing_data, level="info"): + time.sleep(0.01) + + with timer("func2", timing_data, level="debug"): + time.sleep(0.01) + + # Verify separate entries for different functions + assert "func1" in timing_data + assert "func2" in timing_data + assert timing_data["func1"]["count"] == 1 + assert timing_data["func2"]["count"] == 1 + + +class TestSimpleTimerWrap: + """Test cases for the timer_wrap decorator.""" + + def test_timer_wrap_basic_usage(self): + """Test basic functionality of timer_wrap decorator.""" + + class TestClass: + timing_raw: dict = {} + + @timer_wrap(name="test_method", level="info") + def test_method(self): + time.sleep(0.01) + return "result" + + obj = TestClass() + result = obj.test_method() + + # Verify return value + assert result == "result" + + # Verify timing data was recorded + assert hasattr(obj, "timing_raw") + assert "test_method" in obj.timing_raw + assert obj.timing_raw["test_method"]["count"] == 1 + assert obj.timing_raw["test_method"]["total"] > 0 + + def test_timer_wrap_default_name(self): + """Test timer_wrap using default function name.""" + + class TestClass: + timing_raw: dict = {} + + @timer_wrap() + def my_function(self): + time.sleep(0.01) + return 42 + + obj = TestClass() + result = obj.my_function() + + # Verify return value + assert result == 42 + + # Verify timing data was recorded with function name + assert "my_function" in obj.timing_raw + + def test_timer_wrap_multiple_methods(self): + """Test timer_wrap with multiple decorated methods.""" + + class TestClass: + timing_raw: dict = {} + + @timer_wrap(name="method1") + def method1(self): + time.sleep(0.01) + return 1 + + @timer_wrap(name="method2") + def method2(self): + time.sleep(0.02) + return 2 + + obj = TestClass() + obj.method1() + obj.method2() + + # Verify both methods are tracked separately + assert "method1" in obj.timing_raw + assert "method2" in obj.timing_raw + assert obj.timing_raw["method1"]["count"] == 1 + assert obj.timing_raw["method2"]["count"] == 1 + + def test_timer_wrap_multiple_calls(self): + """Test timer_wrap with multiple calls to same method.""" + + class TestClass: + timing_raw: dict = {} + + @timer_wrap(name="repeat_method") + def repeat_method(self): + time.sleep(0.01) + return "done" + + obj = TestClass() + + # Call method multiple times + for _ in range(3): + obj.repeat_method() + + # Verify cumulative statistics + assert obj.timing_raw["repeat_method"]["count"] == 3 + assert obj.timing_raw["repeat_method"]["total"] > 0.03 + + +class TestIntegration: + """Integration tests combining multiple profiler functions.""" + + def test_timer_and_record_timing_compatibility(self): + """Test that timer and record_timing use compatible data structures.""" + timing_data = {} + + # Use timer + with timer("context_timer", timing_data): + time.sleep(0.01) + + # Use record_timing + record_timing("manual_timing", 0.02, timing_data) + + # Verify both use same structure + assert "context_timer" in timing_data + assert "manual_timing" in timing_data + + for key in timing_data: + assert "count" in timing_data[key] + assert "total" in timing_data[key] + assert "avg" in timing_data[key] + + def test_decorator_and_context_manager_compatibility(self): + """Test that decorator and context manager produce compatible data.""" + + class TestClass: + timing_raw: dict = {} + + @timer_wrap(name="decorated_method") + def decorated_method(self): + time.sleep(0.01) + return "done" + + obj = TestClass() + obj.decorated_method() + + timing_data = {} + with timer("context_method", timing_data): + time.sleep(0.01) + + # Verify both produce same data structure + assert "count" in obj.timing_raw["decorated_method"] # type: ignore + assert "total" in obj.timing_raw["decorated_method"] # type: ignore + assert "avg" in obj.timing_raw["decorated_method"] # type: ignore + + assert "count" in timing_data["context_method"] + assert "total" in timing_data["context_method"] + assert "avg" in timing_data["context_method"] + + +if __name__ == "__main__": + pytest.main([__file__]) diff --git a/tests/unit/test_registry.py b/tests/unit/test_registry.py new file mode 100644 index 0000000..041e1c6 --- /dev/null +++ b/tests/unit/test_registry.py @@ -0,0 +1,32 @@ +import pytest + +from rlightning.utils.registry.registry import Registry + + +def test_registry_registers_and_retrieves_default_name(): + registry = Registry("models") + + @registry.register() + class ToyModel: + pass + + assert registry.get("ToyModel") is ToyModel + assert "ToyModel" in registry.module_dict + + +def test_registry_rejects_duplicate_registration(): + registry = Registry("models") + + @registry.register("toy") + class ToyModel: + pass + + with pytest.raises(KeyError, match="already registered"): + registry.register("toy")(ToyModel) + + +def test_registry_raises_for_missing_key(): + registry = Registry("models") + + with pytest.raises(KeyError, match="missing_model is not in the models registry"): + registry.get("missing_model") diff --git a/tests/unit/test_resource_pool.py b/tests/unit/test_resource_pool.py new file mode 100644 index 0000000..e6a97ed --- /dev/null +++ b/tests/unit/test_resource_pool.py @@ -0,0 +1,543 @@ +import pytest + +from rlightning.utils.placement.resource_pool import NodeResource, ResourcePool + +# ============================================================================= +# NodeResource Tests +# ============================================================================= + + +class TestNodeResource: + """Tests for NodeResource.""" + + def test_node_resource_creation(self): + """Test basic NodeResource creation.""" + node = NodeResource( + node_id="node_0", + ip="192.168.1.1", + total_cpus=64, + total_gpus=8, + ) + assert node.node_id == "node_0" + assert node.total_gpus == 8 + assert node.available_gpus == 8 + assert not node.is_empty + + def test_has_resources(self): + """Test resource availability check.""" + node = NodeResource( + node_id="node_0", + ip="192.168.1.1", + total_cpus=64, + total_gpus=8, + ) + assert node.has_resources(gpus=4) + assert node.has_resources(gpus=8) + assert not node.has_resources(gpus=9) + + def test_allocate_basic(self): + """Test basic GPU allocation - now modifies self directly.""" + node = NodeResource( + node_id="node_0", + ip="192.168.1.1", + total_cpus=64, + total_gpus=8, + ) + node.allocate(gpus=4, component_types=["train"]) + + assert node.gpu_cursor == 4 # Cursor advanced + assert node.available_gpus == 4 # Remaining GPUs + assert "train" in node.allocations + assert node.allocations["train"] == [(0, 3)] + + def test_allocate_overlapping_with_consume_false(self): + """Test overlapping allocation using consume=False for colocate mode.""" + node = NodeResource( + node_id="node_0", + ip="192.168.1.1", + total_cpus=64, + total_gpus=8, + ) + # Allocate eval without consuming + node.allocate(gpus=4, component_types=["eval"], consume=False) + # Allocate env on the same GPUs without consuming + node.allocate(gpus=4, component_types=["env"], consume=False) + + assert "eval" in node.allocations + assert "env" in node.allocations + # Both should have the same range (overlapping) + assert node.allocations["eval"] == [(0, 3)] + assert node.allocations["env"] == [(0, 3)] + # Cursor not advanced since consume=False + assert node.gpu_cursor == 0 + assert node.available_gpus == 8 + + def test_allocate_separate_components(self): + """Test separate allocation for multiple components.""" + node = NodeResource( + node_id="node_0", + ip="192.168.1.1", + total_cpus=64, + total_gpus=8, + ) + node.allocate(gpus=[2, 2], component_types=["eval", "env"]) + + assert node.allocations["eval"] == [(0, 1)] + assert node.allocations["env"] == [(2, 3)] + assert node.available_gpus == 4 + + def test_allocate_exceeds_available(self): + """Test allocation exceeding available resources raises error.""" + node = NodeResource( + node_id="node_0", + ip="192.168.1.1", + total_cpus=64, + total_gpus=8, + gpu_cursor=4, # Already allocated 4, only 4 remaining + ) + with pytest.raises(RuntimeError): + node.allocate(gpus=5, component_types=["train"]) + + def test_allocate_zero_gpus(self): + """Test allocation of 0 GPUs logs warning and does nothing.""" + node = NodeResource( + node_id="node_0", + ip="192.168.1.1", + total_cpus=64, + total_gpus=8, + ) + node.allocate(gpus=0) + # Cursor should not advance + assert node.gpu_cursor == 0 + + def test_is_empty(self): + """Test is_empty property.""" + node = NodeResource( + node_id="node_0", + ip="192.168.1.1", + total_cpus=64, + total_gpus=8, + gpu_cursor=8, # All GPUs allocated + ) + assert node.is_empty + + def test_copy(self): + """Test deep copy of NodeResource.""" + node = NodeResource( + node_id="node_0", + ip="192.168.1.1", + total_cpus=64, + total_gpus=8, + allocations={"train": [(0, 3)]}, + gpu_cursor=4, + ) + copied = node.copy() + + assert copied.node_id == node.node_id + assert copied.allocations == node.allocations + assert copied.allocations is not node.allocations # Deep copy + assert copied.gpu_cursor == node.gpu_cursor + + def test_component_types_property(self): + """Test component_types property returns keys from allocations.""" + node = NodeResource( + node_id="node_0", + ip="192.168.1.1", + total_cpus=64, + total_gpus=8, + allocations={"train": [(0, 3)], "buffer": [(4, 4)]}, + ) + assert set(node.component_types) == {"train", "buffer"} + + def test_max_allocated_gpus(self): + """Test max_allocated_gpus tracks the maximum allocation.""" + node = NodeResource( + node_id="node_0", + ip="192.168.1.1", + total_cpus=64, + total_gpus=8, + ) + # First allocation with consume=False + node.allocate(gpus=4, component_types=["train"], consume=False) + assert node.max_allocated_gpus == 4 + assert node.gpu_cursor == 0 # Not consumed + + # Second allocation with consume=False (overlapping) + node.allocate(gpus=6, component_types=["eval"], consume=False) + assert node.max_allocated_gpus == 6 # Updated to max + assert node.gpu_cursor == 0 # Still not consumed + + def test_allocate_consume_true(self): + """Test allocation with consume=True advances cursor.""" + node = NodeResource( + node_id="node_0", + ip="192.168.1.1", + total_cpus=64, + total_gpus=8, + ) + node.allocate(gpus=4, component_types=["train"], consume=True) + assert node.gpu_cursor == 4 + assert node.available_gpus == 4 + assert node.max_allocated_gpus == 4 + + +# ============================================================================= +# ResourcePool Tests +# ============================================================================= + + +class TestResourcePool: + """Tests for ResourcePool.""" + + def test_resource_pool_creation(self): + """Test basic ResourcePool creation with auto-inferred component_types.""" + nodes = [ + NodeResource( + node_id="node_0", + ip="192.168.1.1", + total_cpus=64, + total_gpus=8, + allocations={"train": [(0, 3)]}, + gpu_cursor=4, # 4 GPUs allocated + ) + ] + pool = ResourcePool(name="test_pool", nodes=nodes) + + assert pool.name == "test_pool" + assert pool.num_nodes == 1 + assert pool.total_gpus == 8 + assert "train" in pool.component_types + assert "buffer" in pool.component_types # Auto-added when train exists + + def test_get_component_indices(self): + """Test getting component GPU indices.""" + nodes = [ + NodeResource( + node_id="node_0", + ip="192.168.1.1", + total_cpus=64, + total_gpus=8, + allocations={"train": [(0, 3)]}, + gpu_cursor=4, + ) + ] + pool = ResourcePool(name="test_pool", nodes=nodes) + + indices = pool.get_component_indices("train") + assert indices == "0-3" + + def test_get_component_indices_invalid_component(self): + """Test getting indices for non-existent component.""" + nodes = [ + NodeResource( + node_id="node_0", + ip="192.168.1.1", + total_cpus=64, + total_gpus=8, + allocations={"train": [(0, 3)]}, + gpu_cursor=4, + ) + ] + pool = ResourcePool(name="test_pool", nodes=nodes) + + indices = pool.get_component_indices("eval") + assert indices == "" + + def test_to_dict(self): + """Test serialization to dictionary.""" + nodes = [ + NodeResource( + node_id="node_0", + ip="192.168.1.1", + total_cpus=64, + total_gpus=8, + allocations={"train": [(0, 3)]}, + gpu_cursor=4, + ) + ] + pool = ResourcePool(name="test_pool", nodes=nodes) + + d = pool.to_dict() + assert d["name"] == "test_pool" + assert d["num_node"] == 1 + assert d["num_gpus"] == 8 + assert "train" in d + + def test_to_yaml_dict(self): + """Test YAML-friendly serialization.""" + nodes = [ + NodeResource( + node_id="node_0", + ip="192.168.1.1", + total_cpus=64, + total_gpus=8, + allocations={"train": [(0, 3)]}, + gpu_cursor=4, + ) + ] + pool = ResourcePool(name="test_pool", nodes=nodes) + + yaml_d = pool.to_yaml_dict() + assert yaml_d["name"] == "test_pool" + assert yaml_d["num_node"] == 1 + assert yaml_d["num_gpus"] == 8 + assert "train" in yaml_d + + def test_yaml_dict_heterogeneous_nodes(self): + """Test to_yaml_dict with heterogeneous node GPU counts.""" + nodes = [ + NodeResource( + node_id="node_0", + ip="192.168.1.1", + total_cpus=64, + total_gpus=8, + allocations={"train": [(0, 7)]}, + gpu_cursor=8, + ), + NodeResource( + node_id="node_1", + ip="192.168.1.2", + total_cpus=64, + total_gpus=8, + allocations={"train": [(0, 3)]}, + gpu_cursor=4, + ), + ] + pool = ResourcePool(name="train_pool", nodes=nodes) + + yaml_dict = pool.to_yaml_dict() + assert yaml_dict["name"] == "train_pool" + assert yaml_dict["num_node"] == 2 + assert yaml_dict["num_gpus"] == 8 + assert yaml_dict["train"] == "0-7, 8-11" + + +# ============================================================================= +# ResourcePool YAML Parsing Tests +# ============================================================================= + + +class TestResourcePoolYamlParsing: + """Tests for ResourcePool YAML parsing methods.""" + + def test_parse_index_str_single_range(self): + """Test parsing single range string.""" + result = ResourcePool._parse_index_str("0-7") + assert result == [(0, 7)] + + def test_parse_index_str_multiple_ranges(self): + """Test parsing multiple ranges string.""" + result = ResourcePool._parse_index_str("0-3, 8-11") + assert result == [(0, 3), (8, 11)] + + def test_parse_index_str_single_number(self): + """Test parsing single number.""" + result = ResourcePool._parse_index_str("5") + assert result == [(5, 5)] + + def test_parse_index_str_integer_input(self): + """Test parsing integer input.""" + result = ResourcePool._parse_index_str(5) + assert result == [(5, 5)] + + def test_parse_index_str_none(self): + """Test parsing None input.""" + result = ResourcePool._parse_index_str(None) + assert result == [] + + def test_parse_index_str_empty(self): + """Test parsing empty string.""" + result = ResourcePool._parse_index_str("") + assert result == [] + + def test_parse_index_str_reversed_range(self): + """Test parsing reversed range (end < start).""" + result = ResourcePool._parse_index_str("7-0") + assert result == [(0, 7)] # Should be normalized + + def test_split_global_range_single_node(self): + """Test splitting global range for single node.""" + # Single node with 8 GPUs: offsets = [0, 8] + result = ResourcePool._split_global_range_by_nodes(0, 7, [0, 8]) + assert result == [(0, 0, 7)] # (node_idx, local_start, local_end) + + def test_split_global_range_multi_node(self): + """Test splitting global range across multiple nodes.""" + # Two nodes with 4 GPUs each: offsets = [0, 4, 8] + result = ResourcePool._split_global_range_by_nodes(2, 6, [0, 4, 8]) + # Should split: node 0 gets [2, 3], node 1 gets [0, 2] + assert result == [(0, 2, 3), (1, 0, 2)] + + def test_split_global_range_exact_node_boundary(self): + """Test splitting range that ends exactly at node boundary.""" + # Two nodes with 4 GPUs each: offsets = [0, 4, 8] + result = ResourcePool._split_global_range_by_nodes(0, 3, [0, 4, 8]) + assert result == [(0, 0, 3)] # Only first node + + def test_from_yaml_dict_basic(self): + """Test creating ResourcePool from YAML dict.""" + cluster_nodes = { + "node_0": NodeResource( + node_id="node_0", + ip="192.168.1.1", + total_cpus=64, + total_gpus=8, + ), + "node_1": NodeResource( + node_id="node_1", + ip="192.168.1.2", + total_cpus=64, + total_gpus=8, + ), + } + + pool_cfg = { + "name": "train_pool", + "num_node": 1, + "num_gpus": 8, + "train": "0-7", + } + + pool = ResourcePool.from_yaml_dict(pool_cfg, cluster_nodes) + assert pool.name == "train_pool" + assert pool.num_nodes == 1 + assert "train" in pool.component_types + assert "buffer" in pool.component_types # Auto-added when train exists + + def test_from_yaml_dict_with_node_ids(self): + """Test creating ResourcePool with explicit node_ids.""" + cluster_nodes = { + "node_0": NodeResource( + node_id="node_0", + ip="192.168.1.1", + total_cpus=64, + total_gpus=8, + ), + "node_1": NodeResource( + node_id="node_1", + ip="192.168.1.2", + total_cpus=64, + total_gpus=4, + ), + } + + pool_cfg = { + "name": "train_pool", + "num_node": 2, + "num_gpus": [8, 4], + "node_ids": ["node_0", "node_1"], + "train": "0-11", + } + + pool = ResourcePool.from_yaml_dict(pool_cfg, cluster_nodes) + assert pool.name == "train_pool" + assert pool.num_nodes == 2 + assert pool.node_ids == ["node_0", "node_1"] + + def test_from_yaml_dict_multi_component(self): + """Test creating ResourcePool with multiple components.""" + cluster_nodes = { + "node_0": NodeResource( + node_id="node_0", + ip="192.168.1.1", + total_cpus=64, + total_gpus=8, + ), + } + + pool_cfg = { + "name": "rollout_pool", + "num_node": 1, + "num_gpus": 8, + "eval": "0-3", + "env": "4-7", + } + + pool = ResourcePool.from_yaml_dict(pool_cfg, cluster_nodes) + assert pool.name == "rollout_pool" + assert "eval" in pool.component_types + assert "env" in pool.component_types + + def test_from_yaml_dict_not_enough_nodes(self): + """Test error when not enough nodes available.""" + cluster_nodes = { + "node_0": NodeResource( + node_id="node_0", + ip="192.168.1.1", + total_cpus=64, + total_gpus=8, + ), + } + + pool_cfg = { + "name": "train_pool", + "num_node": 2, # Requires 2 nodes but only 1 available + "num_gpus": 8, + "train": "0-7", + } + + with pytest.raises(ValueError, match="Not enough nodes"): + ResourcePool.from_yaml_dict(pool_cfg, cluster_nodes) + + +# ============================================================================= +# Edge Case Tests +# ============================================================================= + + +class TestEdgeCases: + """Tests for edge cases and boundary conditions.""" + + def test_node_resource_multiple_allocations(self): + """Test multiple sequential allocations on a node - now modifies self directly.""" + node = NodeResource( + node_id="node_0", + ip="192.168.1.1", + total_cpus=64, + total_gpus=8, + ) + + # First allocation - modifies node directly + node.allocate(gpus=4, component_types=["train"]) + assert node.available_gpus == 4 + # Allocations recorded on the node + assert node.allocations["train"] == [(0, 3)] + + # Second allocation - continues from cursor position + node.allocate(gpus=2, component_types=["eval"]) + assert node.available_gpus == 2 + + # Second allocation uses cursor position (continues from 4) + assert node.allocations["eval"] == [(4, 5)] + + # Cursor tracks total allocated + assert node.gpu_cursor == 6 # 4 + 2 + + def test_resource_pool_multi_node(self): + """Test resource pool with multiple nodes.""" + nodes = [ + NodeResource( + node_id="node_0", + ip="192.168.1.1", + total_cpus=64, + total_gpus=8, + allocations={"train": [(0, 3)]}, + gpu_cursor=4, + ), + NodeResource( + node_id="node_1", + ip="192.168.1.2", + total_cpus=64, + total_gpus=8, + allocations={"train": [(0, 3)]}, + gpu_cursor=4, + ), + ] + pool = ResourcePool(name="test_pool", nodes=nodes) + + assert pool.num_nodes == 2 + assert pool.total_gpus == 16 + indices = pool.get_component_indices("train") + # Should have indices from both nodes in pool-global space + assert "0-3" in indices + assert "8-11" in indices diff --git a/tests/unit/test_rollout_buffer.py b/tests/unit/test_rollout_buffer.py new file mode 100644 index 0000000..6461940 --- /dev/null +++ b/tests/unit/test_rollout_buffer.py @@ -0,0 +1,126 @@ +"""Unit tests for RolloutBuffer.""" + +import pytest + +from rlightning.buffer.base_buffer import DataBuffer +from rlightning.buffer.rollout_buffer import RolloutBuffer +from rlightning.buffer.sampler import AllDataSampler, UniformSampler +from rlightning.utils.config import BufferConfig + +def _make_rollout_config(sampler_type: str = "all") -> BufferConfig: + return BufferConfig.from_dict( + { + "type": "RolloutBuffer", + "capacity": 8, + "sampler": {"type": sampler_type}, + "storage": {"type": "unified", "device": "cpu", "unit": "transition"}, + } + ) + + +def test_rollout_buffer_sample_raises_when_batch_too_large(): + """RolloutBuffer should reject batch_size larger than buffer size.""" + buf = RolloutBuffer(_make_rollout_config()) + buf.size = lambda: 2 + + with pytest.raises(ValueError): + buf.sample(batch_size=3) + + +def test_rollout_buffer_sample_allows_equal_batch_size(monkeypatch): + """RolloutBuffer should allow batch_size equal to buffer size.""" + buf = RolloutBuffer(_make_rollout_config()) + buf.size = lambda: 2 + + called = {} + + def _fake_sample(self, batch_size, shuffle=True, drop_last=True): + called["args"] = (batch_size, shuffle, drop_last) + return ["sampled"] + + monkeypatch.setattr(DataBuffer, "sample", _fake_sample) + + cleared = {"flag": False} + + def _fake_clear(): + cleared["flag"] = True + + buf.clear = _fake_clear + + out = buf.sample(batch_size=2, shuffle=False, drop_last=False) + assert out == ["sampled"] + assert called["args"] == (2, False, False) + assert cleared["flag"] is True + + +def test_rollout_buffer_sample_logs_warning_and_clears(caplog, monkeypatch): + """RolloutBuffer should warn on partial sampling and clear after sampling.""" + buf = RolloutBuffer(_make_rollout_config(sampler_type="uniform")) + buf.sampler = UniformSampler() + buf.size = lambda: 5 + + called = {} + + def _fake_sample(self, batch_size, shuffle=True, drop_last=True): + called["args"] = (batch_size, shuffle, drop_last) + return ["sampled"] + + monkeypatch.setattr(DataBuffer, "sample", _fake_sample) + + cleared = {"flag": False} + + def _fake_clear(): + cleared["flag"] = True + + buf.clear = _fake_clear + + with caplog.at_level("WARNING"): + out = buf.sample(batch_size=3, shuffle=False, drop_last=False) + + assert out == ["sampled"] + assert called["args"] == (3, False, False) + assert cleared["flag"] is True + assert "remaining unsampled data will be discarded" in caplog.text + + +def test_rollout_buffer_sample_no_warning_with_all_sampler(caplog, monkeypatch): + """AllDataSampler should avoid partial-sampling warnings.""" + buf = RolloutBuffer(_make_rollout_config(sampler_type="all")) + buf.sampler = AllDataSampler() + buf.size = lambda: 5 + + def _fake_sample(self, batch_size, shuffle=True, drop_last=True): + return ["sampled"] + + monkeypatch.setattr(DataBuffer, "sample", _fake_sample) + + buf.clear = lambda: None + + with caplog.at_level("WARNING"): + _ = buf.sample(batch_size=3) + + assert "remaining unsampled data will be discarded" not in caplog.text + + +def test_rollout_buffer_sample_allows_none_batch_size(caplog, monkeypatch): + """RolloutBuffer should accept batch_size=None without warning.""" + buf = RolloutBuffer(_make_rollout_config(sampler_type="all")) + buf.sampler = AllDataSampler() + buf.size = lambda: 5 + + called = {} + + def _fake_sample(self, batch_size, shuffle=True, drop_last=True): + called["args"] = (batch_size, shuffle, drop_last) + return ["sampled"] + + monkeypatch.setattr(DataBuffer, "sample", _fake_sample) + + buf.clear = lambda: None + + with caplog.at_level("WARNING"): + out = buf.sample(batch_size=None, shuffle=False, drop_last=False) + + assert out == ["sampled"] + assert called["args"] == (None, False, False) + assert "remaining unsampled data will be discarded" not in caplog.text diff --git a/tests/unit/test_router.py b/tests/unit/test_router.py new file mode 100644 index 0000000..e513663 --- /dev/null +++ b/tests/unit/test_router.py @@ -0,0 +1,40 @@ +from rlightning.policy.utils.router import NodeAffinityRouter, SimpleRouter + + +def test_simple_router_balances_assignments_against_current_loads(): + router = SimpleRouter() + current_loads = [2, 0, 1] + + assignments = router.assign(current_loads, num_tasks=4) + + final_loads = [2, 0, 1] + for idx in assignments: + final_loads[idx] += 1 + + assert assignments[0] == 1 + assert max(final_loads) - min(final_loads) <= 1 + + +def test_node_affinity_router_prefers_policies_on_the_same_node(): + router = NodeAffinityRouter( + component_distribution={ + "node-a": {"env": {"ids": [0, 1]}}, + "node-b": {"env": {"ids": [2]}}, + }, + policy_node_ids=["node-a", "node-b", "node-a"], + ) + current_loads = [0, 0, 1] + + assignments = router.assign(current_loads, num_tasks=3, env_ids=["env-0", "env-2", "env-1"]) + + assert assignments == [0, 1, 0] + assert current_loads == [2, 1, 1] + + +def test_node_affinity_router_falls_back_to_simple_routing_without_env_ids(): + router = NodeAffinityRouter(component_distribution={}, policy_node_ids=["node-a", "node-b"]) + current_loads = [1, 0] + + assignments = router.assign(current_loads, num_tasks=2, env_ids=None) + + assert assignments == [1, 0] diff --git a/tests/unit/test_sampler.py b/tests/unit/test_sampler.py new file mode 100644 index 0000000..b7dfc87 --- /dev/null +++ b/tests/unit/test_sampler.py @@ -0,0 +1,60 @@ +import numpy as np +import pytest + +from rlightning.buffer.sampler import AllDataSampler, BatchSampler, UniformSampler + + +def test_uniform_sampler_returns_requested_batch_size_with_valid_indices(): + sampler = UniformSampler() + + indices = sampler.sample(batch_size=20, data_size=5) + + assert indices.shape == (20,) + assert np.all(indices >= 0) + assert np.all(indices < 5) + + +def test_uniform_sampler_can_repeat_indices_when_sampling_with_replacement(): + sampler = UniformSampler() + np.random.seed(0) + + indices = sampler.sample(batch_size=5, data_size=2) + + assert len(indices) == 5 + assert np.all(indices >= 0) + assert np.all(indices < 2) + assert len(set(indices.tolist())) < 5 + + +def test_all_data_sampler_warns_for_partial_batch_and_returns_all_indices(): + sampler = AllDataSampler() + + with pytest.warns(UserWarning, match="sample all data from the buffer"): + indices = sampler.sample(batch_size=2, data_size=4, shuffle=True) + + assert sorted(indices.tolist()) == [0, 1, 2, 3] + + +def test_all_data_sampler_returns_empty_indices_for_empty_input(): + sampler = AllDataSampler() + + indices = sampler.sample(batch_size=0, data_size=0, shuffle=False) + + assert indices.size == 0 + assert indices.tolist() == [] + + +def test_batch_sampler_rejects_oversized_batches(): + sampler = BatchSampler() + + with pytest.raises(ValueError, match="batch_size must be less than or equal to data_size"): + sampler.sample(batch_size=5, data_size=4) + + +def test_batch_sampler_samples_without_replacement(): + sampler = BatchSampler() + + indices = sampler.sample(batch_size=4, data_size=10) + + assert len(indices) == 4 + assert len(set(indices.tolist())) == 4 diff --git a/tests/unit/test_scheduling.py b/tests/unit/test_scheduling.py new file mode 100644 index 0000000..58ab574 --- /dev/null +++ b/tests/unit/test_scheduling.py @@ -0,0 +1,81 @@ +import pytest + +from rlightning.utils.config import MainConfig +from rlightning.utils.placement.scheduling import ComponentScheduling, Scheduling, setup_component_scheduling + + +def test_scheduling_reports_totals_and_dict(): + scheduling = Scheduling(worker_num=4, num_cpus=2, num_gpus=1.0) + + assert scheduling.total_gpus() == 4.0 + assert scheduling.total_cpus() == 8 + assert scheduling.to_dict() == { + "worker_num": 4, + "num_cpus": 2, + "num_gpus": 1.0, + "node_list": None, + } + + +def test_component_scheduling_reports_pool_requirements(): + scheduling = ComponentScheduling( + env_worker=[ + Scheduling(worker_num=2, num_cpus=4, num_gpus=0.5), + Scheduling(worker_num=1, num_cpus=2, num_gpus=0.0), + ], + train_worker=Scheduling(worker_num=2, num_cpus=1, num_gpus=1.0), + eval_worker=Scheduling(worker_num=3, num_cpus=1, num_gpus=0.25), + buffer_worker=Scheduling(worker_num=1, num_cpus=1, num_gpus=0.0), + ) + + assert scheduling.train_pool_requirements() == (2.0, 3) + assert scheduling.rollout_pool_requirements() == (1.75, 13) + assert scheduling.get_component_requirements("env") == (1.0, 10) + + +def test_component_scheduling_rejects_unknown_component_type(): + scheduling = ComponentScheduling() + + with pytest.raises(ValueError, match="Invalid component type"): + scheduling.get_component_requirements("learner") + + +def test_component_scheduling_returns_zero_for_missing_workers(): + scheduling = ComponentScheduling() + + assert scheduling.get_component_requirements("train") == (0, 0) + assert scheduling.get_component_requirements("eval") == (0, 0) + assert scheduling.get_component_requirements("env") == (0, 0) + assert scheduling.get_component_requirements("buffer") == (0, 0) + + +def test_infer_auto_buffer_worker_num_uses_train_gpu_count_and_node_gpu_capacity(): + scheduling = ComponentScheduling( + train_worker=Scheduling(worker_num=5, num_cpus=1, num_gpus=1.0), + buffer_worker=Scheduling(worker_num="auto", num_cpus=1, num_gpus=0.0), + ) + + scheduling.infer_auto_buffer_worker_num( + { + "node_id_to_resources": { + "node-a": {"GPU": 4}, + "node-b": {"GPU": 4}, + } + } + ) + + assert scheduling.buffer_worker.worker_num == 2 + + +def test_setup_component_scheduling_forces_single_buffer_worker_for_unified_storage(make_main_config_dict): + config = MainConfig.from_dict( + make_main_config_dict( + cluster={"buffer_worker_num": 4}, + buffer={"storage": {"type": "unified", "device": "cpu"}}, + ) + ) + + scheduling = setup_component_scheduling(config) + + assert scheduling.buffer_worker.worker_num == 1 + assert config.cluster.buffer_worker_num == 1 diff --git a/tests/unit/test_storage.py b/tests/unit/test_storage.py new file mode 100644 index 0000000..5b26c08 --- /dev/null +++ b/tests/unit/test_storage.py @@ -0,0 +1,896 @@ +"""Unit tests for storage utilities in storage.py.""" + +from __future__ import annotations + +from typing import Any, Dict + +import numpy as np +import pytest +import torch +from tensordict import TensorDict + +from rlightning.buffer.utils.storage import ( + ActiveEpisodeBuffer, + BufferView, + DataContainer, + Storage, +) +from rlightning.types import EnvMeta, EnvRet, PolicyResponse + +def _stack_postprocess(episode: Dict[str, Any]) -> Dict[str, Any]: + """Convert list-based episode fields into tensors for storage.""" + processed: Dict[str, Any] = {} + for key, value in episode.items(): + if isinstance(value, list): + if len(value) == 0: + processed[key] = torch.empty((0,)) + elif isinstance(value[0], torch.Tensor): + processed[key] = torch.stack(value) + else: + processed[key] = torch.tensor(value) + else: + processed[key] = value + return processed + + +def _simple_preprocess( + transition_buffer: Dict[str, Any], + env_ret: EnvRet | None, + policy_resp: PolicyResponse | None, + *_args, + **_kwargs, +) -> Dict[str, Any]: + """Minimal preprocess to assemble a transition dict for tests.""" + if env_ret is not None: + transition_buffer["obs"] = torch.as_tensor(env_ret.observation) + transition_buffer["last_terminated"] = bool(env_ret.last_terminated) + transition_buffer["last_truncated"] = bool(env_ret.last_truncated) + if policy_resp is not None: + transition_buffer["action"] = torch.as_tensor(getattr(policy_resp, "action", 0)) + return transition_buffer + + +def _make_storage(auto_truncate: bool, unit: str = "episode") -> Storage: + return Storage( + capacity=10, + mode="circular", + unit=unit, + env_meta_list=None, + device="cpu", + obs_preprocessor=lambda x: x, + reward_preprocessor=lambda x: x, + env_ret_preprocess_fn=lambda x: x, + policy_resp_preprocess_fn=lambda x: x, + preprocess_fn=_simple_preprocess, + postprocess_fn=_stack_postprocess, + auto_truncate_episode=auto_truncate, + ) + + +def test_buffer_view_len_and_getitem(): + """simple test for BufferView length and indexing.""" + data = TensorDict({"value": torch.tensor([1, 2, 3])}, batch_size=[3]) + view = BufferView(data) + + assert len(view) == 3 + assert view[1]["value"].item() == 2 + + +def test_buffer_view_negative_and_slice_indexing(): + """test for BufferView negative indexing and slicing.""" + data = TensorDict({"value": torch.tensor([5, 6, 7])}, batch_size=[3]) + view = BufferView(data) + + assert view[-1]["value"].item() == 7 + sliced = view[1:] + assert sliced["value"].tolist() == [6, 7] + + +def test_buffer_view_empty_and_full_slice(): + """test for BufferView with empty data and full slice indexing.""" + data = TensorDict({"value": torch.tensor([])}, batch_size=[0]) + view = BufferView(data) + + assert len(view) == 0 + sliced = view[:] + assert sliced["value"].numel() == 0 + + +def test_buffer_view_multidim_batch_len_and_item(): + """test for BufferView with multi-dimensional batch size.""" + data = TensorDict({"value": torch.tensor([[1, 2, 3], [4, 5, 6]])}, batch_size=[2, 3]) + view = BufferView(data) + + assert len(view) == 2 + assert view[1]["value"].tolist() == [4, 5, 6] + + +def test_active_episode_buffer_single_env_pop_done_episode(): + """pop done episode is very import in ActiveEpisodeBuffer. + so we test it as folloing case. we verify only two steps first. + """ + buffer = ActiveEpisodeBuffer( + env_meta_list=None, + auto_truncate_episode=True, + postprocess_fn=_stack_postprocess, + device="cpu", + ) + + buffer.add_transition( + "env-1", + { + "obs": torch.tensor(0), + "last_terminated": False, + "last_truncated": False, + }, + ) + buffer.add_transition( + "env-1", + { + "obs": torch.tensor(1), + "last_terminated": True, + "last_truncated": False, + }, + ) + + episodes = buffer.pop_done_episode("env-1") + + assert len(episodes) == 1 + assert episodes[0]["obs"].tolist() == [0, 1] + assert episodes[0]["last_terminated"].tolist() == [False, True] + assert buffer.pop_done_episode("env-1") == [] + + +def test_pop_done_episode_with_env_id_none(): + """Test env_id=None pops all done episodes from all environments.""" + buffer = ActiveEpisodeBuffer( + env_meta_list=None, + auto_truncate_episode=True, + postprocess_fn=_stack_postprocess, + device="cpu", + ) + + # Add done episode for env-1 + buffer.add_transition( + "env-1", + {"obs": torch.tensor(1), "last_terminated": True, "last_truncated": False}, + ) + # Add done episode for env-2 + buffer.add_transition( + "env-2", + {"obs": torch.tensor(2), "last_terminated": True, "last_truncated": False}, + ) + # Add non-done episode for env-3 + buffer.add_transition( + "env-3", + {"obs": torch.tensor(3), "last_terminated": False, "last_truncated": False}, + ) + + # Pop all done episodes (env_id=None) + episodes = buffer.pop_done_episode(env_id=None) + + assert len(episodes) == 2 + obs_values = sorted([ep["obs"].tolist()[0] for ep in episodes]) + assert obs_values == [1, 2] + + # env-3 should still have data (not done) + assert "env-3" in buffer._buffer + + +def test_pop_done_episode_with_last_truncated(): + """Test last_truncated=True also triggers episode done.""" + buffer = ActiveEpisodeBuffer( + env_meta_list=None, + auto_truncate_episode=True, + postprocess_fn=_stack_postprocess, + device="cpu", + ) + + buffer.add_transition( + "env-1", + {"obs": torch.tensor(0), "last_terminated": False, "last_truncated": False}, + ) + # Use last_truncated instead of last_terminated + buffer.add_transition( + "env-1", + {"obs": torch.tensor(1), "last_terminated": False, "last_truncated": True}, + ) + + episodes = buffer.pop_done_episode("env-1") + + assert len(episodes) == 1 + assert episodes[0]["obs"].tolist() == [0, 1] + assert episodes[0]["last_truncated"].tolist() == [False, True] + + +def test_pop_done_episode_with_sub_env_ids(): + """Test sub-env done flag tracking in multi-env scenarios. + + TODO: Note: When auto_truncate_episode=True and num_envs>1, transitions are split + into sub-env IDs (env-vec/0, env-vec/1). The done flags are tracked per sub-env. + However, pop_done_episode has a limitation: it finds sub-env IDs but pop() + expects parent env_id. This test verifies the done flag tracking behavior. + @yangzhenyu I don't konw whether it will be fixed in future develop branch, so just + check it when you merging. + """ + env_meta_list = [EnvMeta(env_id="env-vec", num_envs=2)] + buffer = ActiveEpisodeBuffer( + env_meta_list=env_meta_list, + auto_truncate_episode=True, + postprocess_fn=_stack_postprocess, + device="cpu", + ) + + # Add transition where only sub-env 0 is done + buffer.add_transition( + "env-vec", + { + "obs": torch.tensor([10, 20]), + "last_terminated": torch.tensor([True, False]), + "last_truncated": torch.tensor([False, False]), + }, + ) + + # Verify sub-envs are stored separately + assert "env-vec/0" in buffer._buffer + assert "env-vec/1" in buffer._buffer + + # Verify done flags are tracked per sub-env + assert buffer._episodes_done_flag["env-vec/0"] is True + assert buffer._episodes_done_flag.get("env-vec/1", False) is False + + # Use pop() with parent env_id to get all sub-env episodes + # (This is the correct way to pop multi-env episodes) + episodes = buffer.pop("env-vec") + + assert len(episodes) == 2 + obs_values = sorted([ep["obs"].tolist()[0] for ep in episodes]) + assert obs_values == [10, 20] + + +def test_pop_done_episode_when_no_done_episodes(): + """Test returns empty list when no episodes are done.""" + buffer = ActiveEpisodeBuffer( + env_meta_list=None, + auto_truncate_episode=True, + postprocess_fn=_stack_postprocess, + device="cpu", + ) + + # Add non-done transitions + buffer.add_transition( + "env-1", + {"obs": torch.tensor(0), "last_terminated": False, "last_truncated": False}, + ) + buffer.add_transition( + "env-1", + {"obs": torch.tensor(1), "last_terminated": False, "last_truncated": False}, + ) + + # No done episodes + episodes = buffer.pop_done_episode("env-1") + assert episodes == [] + + # Also test with env_id=None + episodes = buffer.pop_done_episode(env_id=None) + assert episodes == [] + + # Data should still be in buffer + assert "env-1" in buffer._buffer + assert len(buffer._buffer["env-1"]) == 2 + + +def test_pop_done_episode_multiple_envs_different_states(): + """Test multiple environments with different done states.""" + buffer = ActiveEpisodeBuffer( + env_meta_list=None, + auto_truncate_episode=True, + postprocess_fn=_stack_postprocess, + device="cpu", + ) + + # env-1: 2 transitions, done + buffer.add_transition( + "env-1", + {"obs": torch.tensor(1), "last_terminated": False, "last_truncated": False}, + ) + buffer.add_transition( + "env-1", + {"obs": torch.tensor(2), "last_terminated": True, "last_truncated": False}, + ) + + # env-2: 1 transition, not done + buffer.add_transition( + "env-2", + {"obs": torch.tensor(10), "last_terminated": False, "last_truncated": False}, + ) + + # env-3: 1 transition, done via truncated + buffer.add_transition( + "env-3", + {"obs": torch.tensor(100), "last_terminated": False, "last_truncated": True}, + ) + + # Pop only env-1 + episodes_1 = buffer.pop_done_episode("env-1") + assert len(episodes_1) == 1 + assert episodes_1[0]["obs"].tolist() == [1, 2] + + # Pop env-2 (not done, should be empty) + episodes_2 = buffer.pop_done_episode("env-2") + assert episodes_2 == [] + + # Pop remaining done episodes (env-3) + episodes_remaining = buffer.pop_done_episode(env_id=None) + assert len(episodes_remaining) == 1 + assert episodes_remaining[0]["obs"].tolist() == [100] + + +def test_active_episode_buffer_multi_env_pop_auto_truncate(): + """Test multi-env buffer correctly splits and reassembles transitions when auto_truncate=True.""" + env_meta_list = [EnvMeta(env_id="env-vec", num_envs=2)] + buffer = ActiveEpisodeBuffer( + env_meta_list=env_meta_list, + auto_truncate_episode=True, + postprocess_fn=_stack_postprocess, + device="cpu", + ) + + transition_1 = { + "obs": torch.tensor([1, 2]), + "last_terminated": torch.tensor([False, True]), + "last_truncated": torch.tensor([False, False]), + } + transition_2 = { + "obs": torch.tensor([3, 4]), + "last_terminated": torch.tensor([True, False]), + "last_truncated": torch.tensor([False, False]), + } + + buffer.add_transition("env-vec", transition_1) + buffer.add_transition("env-vec", transition_2) + + episodes = buffer.pop("env-vec") + + # Verify correct number of episodes + assert len(episodes) == 2 + + # Verify obs data correctly split by sub-env + episode_values = [episode["obs"].tolist() for episode in episodes] + assert sorted(episode_values) == sorted([[1, 3], [2, 4]]) + + # Verify other fields also correctly split + term_values = [episode["last_terminated"].tolist() for episode in episodes] + assert sorted(term_values) == sorted([[False, True], [True, False]]) + + # Verify buffer is empty after pop + assert f"env-vec/0" not in buffer._buffer + assert f"env-vec/1" not in buffer._buffer + + +def test_active_episode_buffer_multi_env_pop_no_auto_truncate(): + """Test multi-env pop behavior when auto_truncate is False.""" + env_meta_list = [EnvMeta(env_id="env-vec", num_envs=2)] + buffer = ActiveEpisodeBuffer( + env_meta_list=env_meta_list, + auto_truncate_episode=False, + postprocess_fn=_stack_postprocess, + device="cpu", + ) + + buffer.add_transition( + "env-vec", + { + "obs": torch.tensor([10, 20]), + "last_terminated": torch.tensor([False, False]), + "last_truncated": torch.tensor([False, False]), + }, + ) + buffer.add_transition( + "env-vec", + { + "obs": torch.tensor([30, 40]), + "last_terminated": torch.tensor([False, False]), + "last_truncated": torch.tensor([False, False]), + }, + ) + + episodes = buffer.pop("env-vec") + + assert len(episodes) == 1 + assert episodes[0]["obs"].tolist() == [10, 30, 20, 40] + assert episodes[0]._metadata == {"_num_episodes": 2, "_episode_length": 2} + + +def test_active_episode_buffer_get_num_envs_sub_env_fallback(): + """Test sub-env id resolves to single env.""" + env_meta_list = [EnvMeta(env_id="env-vec", num_envs=2)] + buffer = ActiveEpisodeBuffer( + env_meta_list=env_meta_list, + auto_truncate_episode=True, + postprocess_fn=_stack_postprocess, + device="cpu", + ) + + assert buffer._get_num_envs("env-vec/0") == 1 + + +def test_active_episode_buffer_pop_done_episode_sub_env_ids_returns_only_done(): + """pop_done_episode should handle sub-env ids when auto_truncate=True.""" + env_meta_list = [EnvMeta(env_id="env-vec", num_envs=2)] + buffer = ActiveEpisodeBuffer( + env_meta_list=env_meta_list, + auto_truncate_episode=True, + postprocess_fn=_stack_postprocess, + device="cpu", + ) + + buffer.add_transition( + "env-vec", + { + "obs": torch.tensor([10, 20]), + "last_terminated": torch.tensor([True, False]), + "last_truncated": torch.tensor([False, False]), + }, + ) + + episodes = buffer.pop_done_episode("env-vec") + + assert len(episodes) == 1 + assert episodes[0]["obs"].tolist() == [10] + assert episodes[0]["last_terminated"].tolist() == [True] + + +def test_data_container_transition_getitem_and_use_counter(): + """Test transition getitem updates data_use_counter correctly.""" + container = DataContainer(capacity=3, mode="circular", unit="transition", device="cpu") + container.push({"value": torch.tensor([1, 2, 3])}) + + assert len(container) == 3 + assert container.data_use_counter.tolist() == [0, 0, 0] + + item = container[1] + assert item["value"].item() == 2 + assert container.data_use_counter[1] == 1 + + sliced = container[0:2] + assert sliced["value"].tolist() == [1, 2] + assert container.data_use_counter[0] == 1 + assert container.data_use_counter[1] == 2 + + last = container[-1] + assert last["value"].item() == 3 + assert container.data_use_counter[2] == 1 + + items = container[np.array([0, 2])] + assert items["value"].tolist() == [1, 3] + assert container.data_use_counter[0] == 2 + assert container.data_use_counter[2] == 2 + + +def test_data_container_transition_negative_and_array_indexing(): + """Test transition negative and array indexing paths.""" + container = DataContainer(capacity=4, mode="circular", unit="transition", device="cpu") + container.push({"value": torch.tensor([10, 11, 12, 13])}) + + assert container[-1]["value"].item() == 13 + + indices = np.array([0, -1]) + items = container[indices] + assert items["value"].tolist() == [10, 13] + + with pytest.raises(IndexError): + _ = container[-5] + + with pytest.raises(IndexError): + _ = container[np.array([4])] + + +def test_data_container_transition_circular_wrap(): + """Test circular wrap overwrites transition data as expected.""" + container = DataContainer(capacity=5, mode="circular", unit="transition", device="cpu") + container.push({"value": torch.tensor([0, 1, 2])}) + container.push({"value": torch.tensor([3, 4, 5, 6])}) + + assert container.size == 5 + assert container.pointer == 2 + assert container.data_use_counter.tolist() == [0, 0, 0, 0, 0] + assert container.data["value"].tolist() == [5, 6, 2, 3, 4] + + +def test_data_container_transition_getitem_invalid_index_type_raises(): + """Test transition getitem rejects unsupported index types.""" + container = DataContainer(capacity=3, mode="circular", unit="transition", device="cpu") + container.push({"value": torch.tensor([1, 2])}) + + with pytest.raises(IndexError): + _ = container["invalid"] + + +def test_data_container_transition_fixed_overflow_raises(): + """Test fixed mode raises on capacity overflow.""" + container = DataContainer(capacity=3, mode="fixed", unit="transition", device="cpu") + container.push({"value": torch.tensor([1, 2])}) + + with pytest.raises(RuntimeError): + container.push({"value": torch.tensor([3, 4])}) + + +def test_data_container_clear_resets_state(): + """Test clear resets transition container state.""" + container = DataContainer(capacity=4, mode="circular", unit="transition", device="cpu") + container.push({"value": torch.tensor([1, 2])}) + container.clear() + + assert container.size == 0 + assert container.pointer == 0 + assert np.all(container.data_use_counter == -1) + + +def test_data_container_transition_push_list_items(): + """Test pushing a list of transition dicts.""" + container = DataContainer(capacity=4, mode="circular", unit="transition", device="cpu") + items = [ + {"value": torch.tensor([1])}, + {"value": torch.tensor([2, 3])}, + ] + + container.push(items) + + assert container.size == 3 + assert container.data["value"][: container.size].tolist() == [1, 2, 3] + + +def test_data_container_transition_push_dict_shape_mismatch_raises(): + """Test push rejects dicts with mismatched shapes.""" + container = DataContainer(capacity=4, mode="circular", unit="transition", device="cpu") + bad = {"a": torch.tensor([1, 2]), "b": torch.tensor([1, 2, 3])} + + with pytest.raises((ValueError, RuntimeError)): + container.push(bad) + + +def test_data_container_episode_len_and_getitem_updates_use_counter(): + """Test episode getitem and use counter updates.""" + container = DataContainer(capacity=8, mode="circular", unit="episode", device="cpu") + ep1 = TensorDict({"value": torch.tensor([1, 2])}, batch_size=[2]) + ep2 = TensorDict({"value": torch.tensor([3, 4, 5])}, batch_size=[3]) + + container.push(ep1) + container.push(ep2) + + assert len(container) == 2 + + ep = container[0] + assert ep["value"].tolist() == [1, 2] + assert container.data_use_counter[:2].tolist() == [1, 1] + + eps = container[0:2] + assert len(eps) == 2 + assert container.data_use_counter[:5].tolist() == [2, 2, 1, 1, 1] + + idx_eps = container[np.array([1, -1])] + assert len(idx_eps) == 2 + assert idx_eps[0]["value"].tolist() == [3, 4, 5] + assert container.data_use_counter[:5].tolist() == [2, 2, 3, 3, 3] + + +def test_data_container_episode_invalid_index_type_raises(): + """Test episode getitem rejects unsupported index types.""" + container = DataContainer(capacity=4, mode="circular", unit="episode", device="cpu") + ep = TensorDict({"value": torch.tensor([1])}, batch_size=[1]) + container.push(ep) + + with pytest.raises(TypeError): + _ = container["invalid"] + + +def test_data_container_episode_check_range_out_of_bounds(): + """Test episode index range checks for out-of-bounds access.""" + container = DataContainer(capacity=4, mode="circular", unit="episode", device="cpu") + ep = TensorDict({"value": torch.tensor([1])}, batch_size=[1]) + container.push(ep) + + with pytest.raises(IndexError): + _ = container[1] + + with pytest.raises(IndexError): + _ = container[np.array([1])] + + with pytest.raises(IndexError): + _ = container[torch.tensor([-2])] + + +def test_data_container_episode_overlaps_cleared_on_overwrite(): + """Test episode overlap cleanup when overwriting. + DataContainer is weak here...""" + container = DataContainer(capacity=6, mode="circular", unit="episode", device="cpu") + ep1 = TensorDict({"value": torch.tensor([1, 2, 3])}, batch_size=[3]) + ep2 = TensorDict({"value": torch.tensor([4, 5])}, batch_size=[2]) + ep3 = TensorDict({"value": torch.tensor([6, 7, 8, 9])}, batch_size=[4]) + + container.push(ep1) + container.push(ep2) + # Mark ep1 as used so we can verify it is cleared on overlap. + _ = container[0] + assert container.data_use_counter[:3].tolist() == [1, 1, 1] + + container.push(ep3) + + assert len(container) == 1 # BUG? Should be 2 episodes? + assert container.data_use_counter[:4].tolist() == [0, 0, 0, 0] + assert container.data_use_counter[4:].tolist() == [-1, -1] + + +def test_data_container_episode_circular_large_episode_keeps_last_capacity(): + """Test large episode keeps last capacity in circular mode.""" + container = DataContainer(capacity=3, mode="circular", unit="episode", device="cpu") + ep = TensorDict({"value": torch.tensor([1, 2, 3, 4])}, batch_size=[4]) + + container.push(ep) + + assert len(container) == 1 + assert container.pointer == 0 + assert container.size == 3 + assert container.data["value"].tolist() == [2, 3, 4] + + +def test_data_container_episode_clear_resets_state(): + """Test clear resets episode container state.""" + container = DataContainer(capacity=4, mode="circular", unit="episode", device="cpu") + ep = TensorDict({"value": torch.tensor([1, 2])}, batch_size=[2]) + container.push(ep) + container.clear() + + assert container.size == 0 + assert container.pointer == 0 + assert len(container.episode_registry) == 0 + assert np.all(container.data_use_counter == -1) + + +def test_storage_add_transition_pushes_done_episode(): + """Test add_transition pushes done episodes and updates stats.""" + storage = _make_storage(auto_truncate=True) + + env_ret_1 = EnvRet( + env_id="env-1", + observation=torch.tensor(0), + last_terminated=False, + last_truncated=False, + info={"episode_info": {"score": 1.0}}, + ) + policy_1 = PolicyResponse(env_id="env-1", action=torch.tensor(10)) + + storage.add_transition(env_ret_1, policy_1) + assert len(storage) == 0 + + env_ret_2 = EnvRet( + env_id="env-1", + observation=torch.tensor(1), + last_terminated=True, + last_truncated=False, + info={"episode_info": {"score": 2.0}}, + ) + policy_2 = PolicyResponse(env_id="env-1", action=torch.tensor(11)) + + storage.add_transition(env_ret_2, policy_2) + + assert len(storage) == 1 + episode = storage[0] + assert episode["obs"].shape[0] == 2 + assert episode["action"].shape[0] == 2 + + +def test_storage_len_get_size_get_data_and_clear(): + """Test length/size getters and clear behavior.""" + storage = _make_storage(auto_truncate=True) + + env_ret = EnvRet( + env_id="env-1", + observation=torch.tensor(0), + last_terminated=True, + last_truncated=False, + ) + policy = PolicyResponse(env_id="env-1", action=torch.tensor(1)) + + storage.add_transition(env_ret, policy) + + assert len(storage) == 1 + assert storage.get_size() == 1 + assert storage.size == 1 + + data = storage.get_data() + assert isinstance(data, TensorDict) + assert len(data) == 1 + + storage.clear() + assert len(storage) == 0 + + +def test_storage_add_transition_mismatched_env_ids_raises(): + """Test mismatched env ids raise ValueError.""" + storage = _make_storage(auto_truncate=True) + env_ret = EnvRet(env_id="env-1", observation=torch.tensor(0)) + policy = PolicyResponse(env_id="env-2", action=torch.tensor(0)) + + with pytest.raises(ValueError): + storage.add_transition(env_ret, policy) + + +def test_storage_add_transition_updates_env_ts(): + """Test add_transition updates env timestamp.""" + storage = _make_storage(auto_truncate=False) + + env_ret = EnvRet( + env_id="env-1", + observation=torch.tensor(0), + last_terminated=False, + last_truncated=False, + ) + policy = PolicyResponse(env_id="env-1", action=torch.tensor(1)) + + before = env_ret.ts_env_sent_ns + storage.add_transition(env_ret, policy) + after = env_ret.ts_env_sent_ns + + assert after >= before # I only figure out this validation way. + + +def test_storage_add_data_async_env_only_updates_and_ts(): + """Test async env-only updates timestamp without storing.""" + storage = _make_storage(auto_truncate=True) + + env_ret = EnvRet( + env_id="env-1", + observation=torch.tensor(0), + last_terminated=True, + last_truncated=False, + ) + + before = env_ret.ts_env_sent_ns + storage.add_data_async(env_ret) + after = env_ret.ts_env_sent_ns + + assert after >= before + assert len(storage) == 0 + + +def test_storage_add_data_async_pushes_on_policy_resp_and_rejects_invalid(): + """Test async flow stores on policy resp and rejects invalid.""" + storage = _make_storage(auto_truncate=True) + + env_ret = EnvRet( + env_id="env-1", + observation=torch.tensor(0), + last_terminated=True, + last_truncated=False, + ) + policy = PolicyResponse(env_id="env-1", action=torch.tensor(5)) + + storage.add_data_async(env_ret) + assert len(storage) == 0 + + storage.add_data_async(policy) + assert len(storage) == 1 + assert storage[0]["obs"].shape[0] == 1 + + with pytest.raises(TypeError): + storage.add_data_async("not-a-valid-item") + + +def test_storage_add_episode_splits_multi_env(): + """Test add_episode splits vectorized episodes.""" + storage = _make_storage(auto_truncate=False, unit="episode") + + episode = { + "obs": torch.tensor([[1, 2], [3, 4], [5, 6]]), + "action": torch.tensor([[7, 8], [9, 10], [11, 12]]), + } + + storage.add_episode(episode, num_envs=2) + + assert len(storage) == 2 + episodes = [storage[0], storage[1]] + obs_values = [ep["obs"].tolist() for ep in episodes] + assert sorted(obs_values) == sorted([[1, 3, 5], [2, 4, 6]]) + + +def test_storage_add_episode_single_env_dict(): + """Test add_episode stores single-env dict.""" + storage = _make_storage(auto_truncate=False, unit="episode") + + episode = { + "obs": torch.tensor([1, 2, 3]), + "action": torch.tensor([4, 5, 6]), + } + + storage.add_episode(episode, num_envs=1) + + assert len(storage) == 1 + assert storage[0]["obs"].tolist() == [1, 2, 3] + + +def test_storage_truncate_one_episode_manual_and_invalid_item(): + """Test manual truncate and invalid item handling.""" + storage = _make_storage(auto_truncate=False) + + env_ret = EnvRet( + env_id="env-1", + observation=torch.tensor(0), + last_terminated=False, + last_truncated=False, + ) + policy = PolicyResponse(env_id="env-1", action=torch.tensor(1)) + storage.add_transition(env_ret, policy) + + storage.truncate_one_episode("env-1") + assert len(storage) == 1 + assert storage[0]["obs"].shape[0] == 1 + + class _NoEnvId: + pass + + with pytest.raises(TypeError): + storage.truncate_one_episode(_NoEnvId()) + + +def test_storage_truncate_one_episode_object_with_env_id(): + """Test truncate_one_episode accepts objects with env_id.""" + storage = _make_storage(auto_truncate=False) + + env_ret = EnvRet( + env_id="env-1", + observation=torch.tensor(0), + last_terminated=False, + last_truncated=False, + ) + policy = PolicyResponse(env_id="env-1", action=torch.tensor(1)) + storage.add_transition(env_ret, policy) + + class _HasEnvId: + env_id = "env-1" + + storage.truncate_one_episode(_HasEnvId()) + + assert len(storage) == 1 + + +def test_storage_truncate_episodes_calls_truncate_one_episode(): + """Test truncate_episodes truncates each env id.""" + storage = _make_storage(auto_truncate=False) + + env_ret = EnvRet( + env_id="env-1", + observation=torch.tensor(0), + last_terminated=False, + last_truncated=False, + ) + policy = PolicyResponse(env_id="env-1", action=torch.tensor(1)) + storage.add_transition(env_ret, policy) + + storage.truncate_episodes(["env-1"]) # only one but it is a list + + assert len(storage) == 1 + + +def test_storage_print_timing_summary_resets(): + """Test print_timing_summary reset behavior.""" + storage = _make_storage(auto_truncate=False) + storage.timing_raw = {"x": {"count": 1, "total": 1.0, "avg": 1.0}} + + storage.print_timing_summary(reset=True) + + assert storage.timing_raw == {} + + +def test_storage_print_timing_summary_logs_contents(caplog): + """Test print_timing_summary logs timing details.""" + storage = _make_storage(auto_truncate=False) + storage.timing_raw = {"transition_pair_to_buffer": {"count": 2, "total": 1.5, "avg": 0.75}} + + with caplog.at_level("DEBUG"): + storage.print_timing_summary(reset=False) + + assert "Buffer storage timing:" in caplog.text + assert "transition_pair_to_buffer" in caplog.text + assert "count=2" in caplog.text diff --git a/tests/unit/test_sync_rl_engine.py b/tests/unit/test_sync_rl_engine.py new file mode 100644 index 0000000..cc0ff6d --- /dev/null +++ b/tests/unit/test_sync_rl_engine.py @@ -0,0 +1,273 @@ +import builtins +import os +from pathlib import Path +from types import SimpleNamespace +from typing import Dict, Tuple +from unittest import mock + +import gymnasium as gym +import numpy as np +import pytest +import torch +from hydra import compose, initialize_config_dir +from omegaconf import DictConfig, OmegaConf + +from rlightning.engine.sync_rl_engine import SyncRLEngine +from rlightning.env import EnvMeta +from rlightning.policy.base_policy import BasePolicy, PolicyRole +from rlightning.policy.policy_group import PolicyGroup +from rlightning.types import BatchedData, EnvRet +from rlightning.utils.config import ClusterConfig, PolicyConfig, TrainConfig +from rlightning.utils.registry import POLICIES + +# Resolve repo root from tests/unit/test_sync_rl_engine.py. +EXAMPLES_DIR = Path(__file__).resolve().parents[2] / "examples" + + +def policy_intialization_entrypoint(cfg: DictConfig, as_remote: bool, role: str) -> BasePolicy: + """Help function to initialize BasePolicy from config. + + This function can be used to test whether the configuration is valid. + + Args: + cfg (DictConfig): Configuration of BasePolicy. + as_remote (bool): Enable remote mode or not. + role (str): Indicate current policy is for 'train' or 'test'. + + Raises: + NotImplementedError: Remote mode is not support yet. + + Returns: + BasePolicy: An instance of BasePolicy. + """ + + print(f"--- Complete Config ---\n{cfg}") + + policy_cfg = PolicyConfig.from_omegaconf(cfg) + + policy_cfg.as_remote = as_remote + policy_cls = POLICIES.get(policy_cfg.type) + + if as_remote: + + class DistributedPolicyCls(policy_cls, DistributedMixin): + pass + + policy_cls = DistributedPolicyCls + + if as_remote: + policy_name = f"policy-{role}-{policy_cfg.type}-test" + + raise NotImplementedError + else: + policy = policy_cls(policy_cfg, role_type=PolicyRole(role)) + + return policy + + +def make_engine_with_mocks(**cfg_overrides) -> Tuple[SyncRLEngine, Dict]: + """Create a SyncRLEngine instance with mocked components for testing.""" + + # Create instance without calling __init__ + engine = SyncRLEngine.__new__(SyncRLEngine) + + # Minimal config with nested train and policy attributes + train_cfg = TrainConfig( + max_epochs=2, + max_rollout_steps=2, + batch_size=10, + lr=0.0003, + mode="sync", + parallel=None, + eval_interval=2, + save_interval=5, + save_dir="./save_dir", + ) + config = SimpleNamespace(train=train_cfg) + + if cfg_overrides.get("train"): + for k, v in cfg_overrides["train"].items(): + setattr(config.train, k, v) + + if cfg_overrides.get("policy"): + config.policy = cfg_overrides["policy"] + else: + config.policy = mock.Mock() + + config.cluster = ClusterConfig() + + # Create mocks for components + env_group = mock.Mock() + + # auto_reset context manager mock + auto_reset_cm = mock.MagicMock() + auto_reset_cm.__enter__.return_value = env_group + auto_reset_cm.__exit__.return_value = False + env_group.auto_reset.return_value = auto_reset_cm + + # init returns metadata + env_meta = EnvMeta( + env_id=None, + action_space=gym.spaces.Box(low=-np.inf, high=np.inf, shape=(4,), dtype=np.float32), + observation_space=gym.spaces.Box(low=-np.inf, high=np.inf, shape=(16,)), + num_envs=2, + ) + env_group.init.return_value = [env_meta] + + batched_action_space = gym.spaces.Box( + low=-np.inf, + high=np.inf, + shape=(env_meta.num_envs,) + env_meta.action_space.shape, + dtype=np.float32, + ) + batched_observation_space = gym.spaces.Box( + low=-np.inf, + high=np.inf, + shape=(env_meta.num_envs,) + env_meta.observation_space.shape, + dtype=np.float32, + ) + + # reset/step yield a batched_env_ret and truncations + batched_env_ret = BatchedData( + ids=[0], + data=[ + EnvRet( + env_id=0, + observation=torch.tensor(batched_observation_space.sample()), + last_reward=torch.zeros(env_meta.num_envs), + last_terminated=torch.zeros(env_meta.num_envs, dtype=torch.bool), + last_truncated=torch.zeros(env_meta.num_envs, dtype=torch.bool), + info={}, + ) + ], + ) + truncations = None + env_group.reset.return_value = (batched_env_ret, truncations) + env_group.step.return_value = (batched_env_ret, truncations) + # get_env_stats returns a dict for logging + env_group.get_env_stats.return_value = {"reward_mean": 1.2345} + + buffer = mock.MagicMock() + buffer.init.return_value = None + buffer.add_batched_transition.return_value = None + buffer.truncate_episodes.return_value = None + # buffer.sample returns a dummy dataset + sample_data = {"data": "sample"} + buffer.sample.return_value = sample_data + buffer.clear.return_value = None + # Provide sampler attribute used in train + buffer.sampler = SimpleNamespace(replacement=True) + + # Provide a default policy_group mock so rollout/train won’t fail + batched_policy_resp = mock.Mock() + policy_group = mock.Mock() + policy_group.rollout_batch.return_value = batched_policy_resp + policy_group.send_weights.return_value = None + policy_group.notify_update_weights.return_value = None + policy_group.update_dataset.return_value = None + policy_group.train.return_value = None + + # timing_raw used by profiler.timer but we can leave as empty dict + engine.timing_raw = {} + + # attach mocks and config to engine + engine.env_group = env_group + engine.policy_group = policy_group + engine.buffer = buffer + engine.config = config + engine.epoch = 0 + + return engine, { + "env_meta": env_meta, + "batched_env_ret": batched_env_ret, + # "batched_policy_resp": batched_policy_resp, + "sample_data": sample_data, + } + + +def test_warm_up_calls_sequence(): + engine, ctx = make_engine_with_mocks() + + # Mock policy_group with expected methods + policy_group = mock.Mock() + episode_meta = {"episode": "meta"} + policy_group.init_eval.return_value = ctx["env_meta"] + policy_group.init_train.return_value = None + + # batched policy response mock with ids method + batched_policy_resp = mock.Mock() + batched_policy_resp.ids.return_value = [1, 2] + policy_group.rollout_batch.return_value = batched_policy_resp + + policy_group.update_dataset.return_value = None + policy_group.train.return_value = None + policy_group.send_weights.return_value = None + policy_group.notify_update_weights.return_value = None + policy_group.train_list = [mock.Mock()] + + engine.policy_group = policy_group + + # Call warm_up and assert all expected interactions occur + engine.env_meta_list = [ctx["env_meta"]] + engine.warm_up() + + # truncate_episodes called with ids() of last batched_policy_resp + engine.buffer.truncate_episodes.assert_called_once_with(batched_policy_resp.ids()) + + # dummy run train: sample and update_dataset and train called + engine.buffer.sample.assert_called_once_with(batch_size=engine.config.train.batch_size) # defined in warm_up + engine.policy_group.update_dataset.assert_called_once_with(ctx["sample_data"]) + engine.policy_group.train.assert_called_once() + + # buffer cleared + engine.buffer.clear.assert_called_once() + + +def test_rollout_calls_add_batched_transition_and_get_env_stats(monkeypatch, capsys): + # Use small max_rollout_steps to verify loop behavior + engine, ctx = make_engine_with_mocks() + engine.config.train.max_rollout_steps = 3 + + # Ensure auto_reset context manager returns env_group (mock already configured) + engine.rollout(obj_set="train", prefix="rollout/") + + # Expected number of calls: max_rollout_steps + 1 iterations + expected_calls = engine.config.train.max_rollout_steps + 1 + assert engine.buffer.add_batched_transition.call_count == expected_calls + + # get_env_stats should be called once with reset=True and prefix + engine.env_group.get_env_stats.assert_called_once_with(reset=True) + + +def test_update_weights_calls_policy_group(): + engine, ctx = make_engine_with_mocks() + engine.sync_weights() + + engine.policy_group.sync_weights.assert_called_once_with() + + +def test_run_invokes_flow_and_update_weights_calls(monkeypatch): + engine, ctx = make_engine_with_mocks() + # Replace rollout/train/update_weights with mocks to track calls + engine.rollout = mock.Mock() + engine.update_dataset = mock.Mock() + engine.train = mock.Mock() + engine.sync_weights = mock.Mock() + + # Ensure epochs small and intervals set so update branches run + engine.config.train.max_epochs = 3 + engine.config.train.save_interval = 1 + engine.config.train.eval_interval = 1 + + # Disable verbose/progress to avoid side-effects + os.environ["RLIGHTNING_VERBOSE"] = "0" + + engine.run() + + # rollout should be called at least once per epoch for training and evals as configured + assert engine.rollout.call_count >= engine.config.train.max_epochs + # train should be called once per epoch + assert engine.train.call_count == engine.config.train.max_epochs + # sync_weights should be called once per epoch + assert engine.sync_weights.call_count == engine.config.train.max_epochs + # restore flags if needed (no-op in this test scope) diff --git a/tests/unit/test_table.py b/tests/unit/test_table.py new file mode 100644 index 0000000..537f67c --- /dev/null +++ b/tests/unit/test_table.py @@ -0,0 +1,52 @@ +"""Unit tests for EpisodeTable.""" + +import pytest + +from rlightning.buffer.utils.table import EpisodeTable + +# TODO: add test for new functions + + +def test_episode_table_init_requires_positive_storages(): + """EpisodeTable should require num_storages >= 1.""" + with pytest.raises(ValueError): + EpisodeTable(num_storages=0) + + +def test_episode_table_register_envs_balances_round_robin(): + """EpisodeTable should assign envs to lowest-load shard with RR tie-break.""" + table = EpisodeTable(num_storages=2) + table.register_envs(["env-a", "env-b", "env-c"]) + + assert table.get_storage_idx_for_env("env-a") == 0 + assert table.get_storage_idx_for_env("env-b") == 1 + assert table.get_storage_idx_for_env("env-c") == 0 + + +def test_episode_table_get_storage_idx_is_stable_for_existing_env(): + """get_storage_idx should return the same shard for an existing env.""" + table = EpisodeTable(num_storages=3) + idx_first = table.get_storage_idx_for_env("env-x") + idx_second = table.get_storage_idx_for_env("env-x") + + assert idx_first == idx_second + assert table._storage_env_count[idx_first] == 1 + + +def test_episode_table_envs_for_storage_lists_envs(): + """envs_for_storage should list env ids assigned to a shard.""" + table = EpisodeTable(num_storages=2, env_ids=["env-1", "env-2", "env-3"]) + + envs0 = table.get_envs_for_storage(0) + envs1 = table.get_envs_for_storage(1) + + assert set(envs0) == {"env-1", "env-3"} + assert set(envs1) == {"env-2"} + + +def test_episode_table_envs_for_storage_validates_index(): + """envs_for_storage should validate storage index.""" + table = EpisodeTable(num_storages=2) + + with pytest.raises(IndexError): + _ = table.get_envs_for_storage(2) diff --git a/tests/unit/test_tensor_utils.py b/tests/unit/test_tensor_utils.py new file mode 100644 index 0000000..4eaaac7 --- /dev/null +++ b/tests/unit/test_tensor_utils.py @@ -0,0 +1,64 @@ +import numpy as np +import pytest +import torch + +from rlightning.utils.utils import InternalFlag, to_device, to_numpy, torch_dtype_from_precision + + +@pytest.mark.parametrize( + ("precision", "expected_dtype"), + [ + ("bf16", torch.bfloat16), + ("bf16-mixed", torch.bfloat16), + (16, torch.float16), + ("16", torch.float16), + ("fp16", torch.float16), + ("16-mixed", torch.float16), + (32, torch.float32), + ("32", torch.float32), + ("32-true", torch.float32), + (None, None), + ], +) +def test_torch_dtype_from_precision_variants(precision, expected_dtype): + assert torch_dtype_from_precision(precision) is expected_dtype + + +def test_torch_dtype_from_precision_rejects_unknown_value(): + with pytest.raises(ValueError, match="Could not parse the precision"): + torch_dtype_from_precision("fp8") + + +def test_bfloat16_numpy_round_trip_preserves_values(): + source = {"weights": torch.tensor([1.5, -2.0], dtype=torch.bfloat16)} + + numpy_data = to_numpy(source) + restored = to_device(numpy_data, "cpu") + + assert isinstance(numpy_data["weights"], np.ndarray) + assert numpy_data["weights"].dtype == np.uint16 + assert restored["weights"].dtype == torch.bfloat16 + assert torch.equal(restored["weights"], source["weights"]) + + +def test_to_numpy_rejects_raw_uint16_tensors(): + with pytest.raises(ValueError, match="haven't support converting uint16 tensor"): + to_numpy(torch.tensor([1, 2], dtype=torch.uint16)) + + +def test_internal_flag_get_env_vars_reflects_environment(monkeypatch): + monkeypatch.setenv("RLIGHTNING_DEBUG", "1") + monkeypatch.setenv("RLIGHTNING_VERBOSE", "0") + monkeypatch.setenv("RLIGHTNING_REMOTE_TRAIN", "1") + monkeypatch.setenv("RLIGHTNING_REMOTE_EVAL", "0") + monkeypatch.setenv("RLIGHTNING_REMOTE_STORAGE", "1") + monkeypatch.setenv("RLIGHTNING_REMOTE_ENV", "0") + + assert InternalFlag.get_env_vars() == { + "RLIGHTNING_DEBUG": "1", + "RLIGHTNING_VERBOSE": "0", + "RLIGHTNING_REMOTE_TRAIN": "1", + "RLIGHTNING_REMOTE_EVAL": "0", + "RLIGHTNING_REMOTE_STORAGE": "1", + "RLIGHTNING_REMOTE_ENV": "0", + } diff --git a/tests/unit/test_vision_models.py b/tests/unit/test_vision_models.py new file mode 100644 index 0000000..70c5322 --- /dev/null +++ b/tests/unit/test_vision_models.py @@ -0,0 +1,21 @@ +from typing import Tuple + +import pytest +import torch + +from rlightning.models.vision.cnn import NatureCNN + +@pytest.mark.parametrize("image_shape", [(84, 84), (210, 160)]) +@pytest.mark.parametrize("batch_size", [1, 4]) +@pytest.mark.parametrize("image_format", ["HWC", "CHW"]) +def test_any_shape_as_input(image_shape: Tuple[int, int], batch_size: int, image_format: str): + model = NatureCNN(image_format=image_format) + + if image_format == "HWC": + observation = torch.rand((batch_size,) + image_shape + (model.in_channels,)) + else: + observation = torch.rand((batch_size, model.in_channels) + image_shape) + + features = model(obs=observation) + + assert features.shape == (batch_size, model.out_feature_dim) diff --git a/tests/unit/test_vtrace.py b/tests/unit/test_vtrace.py new file mode 100644 index 0000000..9f8e3f5 --- /dev/null +++ b/tests/unit/test_vtrace.py @@ -0,0 +1,53 @@ +import torch + +from rlightning.policy.utils.vtrace import batch_step_correction, vtrace_correction + + +def test_vtrace_correction_reduces_to_clipped_td_error_when_gamma_is_zero(): + rewards = torch.tensor([1.0, 2.0]) + values = torch.tensor([0.5, 1.5]) + next_values = torch.tensor([10.0, 20.0]) + dones = torch.tensor([False, True]) + log_rhos = torch.log(torch.tensor([2.0, 0.5])) + + vs, advantages = vtrace_correction( + rewards=rewards, + values=values, + next_values=next_values, + dones=dones, + log_rhos=log_rhos, + gamma=0.0, + rho_bar=1.0, + c_bar=1.0, + ) + + clipped_rhos = torch.tensor([1.0, 0.5]) + expected = clipped_rhos * (rewards - values) + + assert torch.allclose(vs, values + expected) + assert torch.allclose(advantages, expected) + + +def test_batch_step_correction_reduces_to_clipped_td_error_when_gamma_is_zero(): + rewards = torch.tensor([[1.0, 2.0]]) + values = torch.tensor([[0.5, 1.5]]) + next_values = torch.tensor([[10.0, 20.0]]) + log_rhos = torch.log(torch.tensor([[2.0, 0.5]])) + dones = torch.tensor([[False, True]]) + + target_values, advantages = batch_step_correction( + rewards=rewards, + values=values, + next_values=next_values, + log_rhos=log_rhos, + dones=dones, + gamma=0.0, + rho_bar=1.0, + c_bar=1.0, + ) + + clipped_rhos = torch.tensor([[1.0, 0.5]]) + expected = clipped_rhos * (rewards - values) + + assert torch.allclose(target_values, values + expected) + assert torch.allclose(advantages, expected) diff --git a/third_party/rsl_rl b/third_party/rsl_rl new file mode 160000 index 0000000..f7d99e8 --- /dev/null +++ b/third_party/rsl_rl @@ -0,0 +1 @@ +Subproject commit f7d99e8e6bb12762d20daadf5d3b5e143b1200b3