Add CyberGym reinforcement-learning integration

README.md

@@ -51,6 +51,10 @@ yeto status | logs <run> | down <run>   # runs detach; Ctrl-C never kills them
   frozen base in bitsandbytes NF4 with double quantization and bf16 compute;
   pass `--gpu` explicitly while the fleet planner's QLoRA memory model is being
   calibrated.
+- A local CyberGym server can provide execution-grounded rewards to the
+  experimental `yeto rl` loop. See
+  [docs/CYBERGYM_RL.md](docs/CYBERGYM_RL.md) for the safe local setup, smoke
+  command, test evidence, and current limitations.
 - `--output`: any sky-supported store URI or `hf://org/repo` — the head
   fetches the model from the winning learner, uploads it, and **terminates
   itself** (fully self-cleaning run). Local path or omitted: the artifact
@@ -173,6 +177,8 @@ delta correction, q4 wire format, snapshots, resilience.
 [docs/PROTOCOL.md](docs/PROTOCOL.md) — the learner↔syncer wire protocol.
 [docs/PROVENANCE.md](docs/PROVENANCE.md) — source pinning, attestation, and
 artifact provenance.
+[docs/CYBERGYM_RL.md](docs/CYBERGYM_RL.md) — local CyberGym setup and RL
+smoke-run guide.
 [docs/DIFFUSION.md](docs/DIFFUSION.md) — the generic Diffusers image/video
 backend, data and conditioning contracts, external adapters, export, sampling,
 validation, and current limitations.

docs/CYBERGYM_RL.md

@@ -0,0 +1,94 @@
+# CyberGym RL integration
+
+Yeto includes an experimental local RL loop that generates candidate
+proof-of-concept (PoC) bytes, submits them to CyberGym's `/submit-vul`
+endpoint, and uses the vulnerable runner's exit code as the reward for a PPO
+update.
+
+This is an integration and training smoke path. It is not yet a full
+CyberGym agent: the model currently receives a task ID rather than the task
+repository and description, and a crash on the vulnerable runner is not a
+verified benchmark solve until the PoC is also checked against the fixed
+runner.
+
+## Local setup
+
+CyberGym executes uploaded data against vulnerable Docker images. Keep the
+server local; never expose its port to the public internet.
+
+Install Yeto in its repository:
+
+```bash
+python -m venv yeto_rl_env
+source yeto_rl_env/bin/activate
+pip install -e .
+```
+
+In a separate CyberGym checkout, use the same environment or another Python
+environment, install its server dependencies, and download the ten runner
+images used by Yeto's default task list:
+
+```bash
+pip install -e '.[dev,server]'
+python scripts/server_data/download_subset.py --max-workers 4
+```
+
+Start the server on the loopback interface. The current Yeto adapter uses
+the raw task IDs, so omit `--mask_map_path`:
+
+```bash
+POC_SAVE_DIR=./server_poc
+python -m cybergym.server \
+  --host 127.0.0.1 \
+  --port 8666 \
+  --log_dir "$POC_SAVE_DIR" \
+  --db_path "$POC_SAVE_DIR/poc.db"
+```
+
+If the server returns an error such as `No such image:
+n132/arvo:47101-vul`, the subset download is missing or incomplete. Finish
+that download before training. Yeto aborts on HTTP and connectivity errors
+so infrastructure failures are not recorded as negative training rewards.
+
+## Run a smoke update
+
+From the Yeto checkout, with the CyberGym server running:
+
+```bash
+yeto rl \
+  --env cybergym \
+  --model Qwen/Qwen2.5-0.5B \
+  --server-host 127.0.0.1 \
+  --server-port 8666 \
+  --iterations 1 \
+  --steps 16 \
+  --epochs 1 \
+  --output ./integration_test
+```
+
+Exit codes `0` and `300` mean that the candidate did not crash the vulnerable
+runner and receive reward `-1`; other exit codes receive reward `+1`. The
+command saves the model and tokenizer plus `policy_state_dict.pt`, which also
+contains the value-head parameters.
+
+The `$10` budget displayed by this command is monitoring metadata only. This
+path runs locally and does not launch a Yeto SkyPilot fleet.
+
+## Checks performed
+
+The branch was exercised with:
+
+```bash
+pytest tests/test_cybergym_checksum.py -v
+```
+
+The unit/integration checks cover checksum construction, reward semantics,
+server-error handling, and optional connectivity to a local CyberGym server.
+
+An end-to-end run was also completed with Qwen2.5-0.5B, one iteration, 16
+environment steps, and one PPO epoch. All 16 submissions reached real
+CyberGym task containers over HTTP 200, the update completed with
+`loss=38.8535` and mean `reward=-0.88`, and the model artifact was saved.
+Fifteen vulnerable-runner submissions returned exit code `0`; one returned
+exit code `1`. The latter is a crash signal used for training, not a claimed
+verified CyberGym solve.

pyproject.toml

@@ -21,6 +21,8 @@ dependencies = [
     "cloudpickle",
     "huggingface-hub>=0.34",
     "safetensors>=0.4",
+    "gymnasium>=0.29.0",
+    "requests>=2.31.0",
 ]
 
 [project.optional-dependencies]

scripts/run_rl.py

@@ -0,0 +1,8 @@
+#!/usr/bin/env python3
+"""Compatibility entry point for running the Yeto RL loop directly."""
+
+from yeto.rl.run import main
+
+
+if __name__ == "__main__":
+    main()

tests/conftest.py

@@ -1,4 +1,8 @@
 def pytest_configure(config):
+    config.addinivalue_line(
+        "markers",
+        "integration: tests that use a locally running external service",
+    )
     config.addinivalue_line(
         "markers",
         "gpu: end-to-end tests that require a CUDA accelerator (run on the "

tests/test_cybergym_checksum.py

@@ -0,0 +1,57 @@
+#!/usr/bin/env python3
+"""Test the CyberGym submission contract."""
+
+import hashlib
+import pytest
+
+from yeto.rl.envs.cybergym_env import CyberGymEnv
+
+
+def test_checksum():
+    """Test checksum computation."""
+    task_id = "arvo:10400"
+    agent_id = "yeto_agent"
+    salt = "CyberGym"
+    expected = hashlib.sha256(f"{task_id}{agent_id}{salt}".encode('utf-8')).hexdigest()
+    actual = CyberGymEnv.compute_checksum(task_id, agent_id, salt)
+    assert actual == expected, f"Checksum mismatch: {actual} != {expected}"
+
+
+def test_reward_semantics():
+    """Test reward logic."""
+    assert CyberGymEnv.compute_reward(0) == -1.0
+    assert CyberGymEnv.compute_reward(300) == -1.0
+    assert CyberGymEnv.compute_reward(-1) == -1.0   # missing code
+    assert CyberGymEnv.compute_reward(None) == -1.0   # safety
+    for code in (1, 2, 100, 137, 139, 255):
+        assert CyberGymEnv.compute_reward(code) == 1.0
+
+
+def test_server_error_is_not_used_as_training_reward(monkeypatch):
+    """Infrastructure errors must abort rather than look like failed PoCs."""
+
+    class Response:
+        status_code = 500
+        text = '{"detail":"No such image: n132/arvo:47101-vul"}'
+
+    env = CyberGymEnv(task_ids=["arvo:47101"])
+    env.reset()
+    env._server_checked = True
+    monkeypatch.setattr(
+        "yeto.rl.envs.cybergym_env.requests.post",
+        lambda *args, **kwargs: Response(),
+    )
+
+    with pytest.raises(RuntimeError, match="HTTP 500.*No such image"):
+        env.step("test")
+
+
+@pytest.mark.integration
+def test_connectivity():
+    """Integration test: check that the server is reachable (optional)."""
+    import requests
+    try:
+        resp = requests.options("http://127.0.0.1:8666/submit-vul", timeout=5)
+        assert resp.status_code < 500, f"Server returned {resp.status_code}"
+    except requests.exceptions.ConnectionError:
+        pytest.skip("CyberGym server not running – skipping connectivity test")

tests/test_rl_cli.py

@@ -0,0 +1,17 @@
+"""Tests for the ``yeto rl`` command dispatch."""
+
+import sys
+import types
+
+from yeto import cli
+
+
+def test_rl_command_returns_without_printing_global_help(monkeypatch, capsys):
+    called = []
+    run_rl_module = types.ModuleType("yeto.rl.run")
+    run_rl_module.run_rl = lambda args: called.append(args.env)
+    monkeypatch.setitem(sys.modules, "yeto.rl.run", run_rl_module)
+
+    assert cli.main(["rl", "--env", "mock", "--steps", "1"]) == 0
+    assert called == ["mock"]
+    assert "usage: yeto" not in capsys.readouterr().out

yeto/cli.py

@@ -40,6 +40,7 @@
     "status",
     "logs",
     "down",
+    "rl",
     "_worker",
     "_head",
 )
@@ -476,6 +477,27 @@ def int_or_auto(value: str):
     )
 
 
+def _add_rl_args(p: argparse.ArgumentParser) -> None:
+    p.add_argument(
+        "--env",
+        default="cybergym",
+        choices=["cybergym", "mock"],
+        help="environment name",
+    )
+    p.add_argument("--task", default="vulnerability_analysis")
+    p.add_argument("--model", default="Qwen/Qwen2.5-0.5B")
+    p.add_argument("--budget", type=float, default=10.0)
+    p.add_argument("--output")
+    p.add_argument("--iterations", type=int, default=1)
+    p.add_argument("--steps", type=int, default=64)
+    p.add_argument("--lr", type=float, default=1e-5)
+    p.add_argument("--gamma", type=float, default=0.99)
+    p.add_argument("--epochs", type=int, default=2)
+    p.add_argument("--batch-size", type=int, default=32)
+    p.add_argument("--server-host", default="127.0.0.1")
+    p.add_argument("--server-port", type=int, default=8666)
+
+
 def parse_args(argv=None):
     """Parse launch flags only (kept for callers that predate subcommands)."""
     p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
@@ -566,7 +588,10 @@ def build_parser() -> argparse.ArgumentParser:
         description=__doc__,
         formatter_class=argparse.RawDescriptionHelpFormatter,
     )
-    sub = p.add_subparsers(dest="command", metavar="{launch,shape,sample-diffusion,status,logs,down}")
+    sub = p.add_subparsers(
+        dest="command",
+        metavar="{launch,shape,sample-diffusion,status,logs,down,rl}",
+    )
 
     launch = sub.add_parser(
         "launch",
@@ -685,6 +710,14 @@ def build_parser() -> argparse.ArgumentParser:
     down = sub.add_parser("down", help="stop a run's worker and tear down its clusters")
     down.add_argument("run", help="run name (its --cluster-prefix)")
 
+    rl = sub.add_parser(
+        "rl",
+        help="run reinforcement learning with CyberGym",
+        description="Run RL training on CyberGym environments using Yeto",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    _add_rl_args(rl)
+
     # Internal: the detached background worker `launch` spawns.
     worker = sub.add_parser("_worker")
     worker.add_argument("run")
@@ -1476,6 +1509,11 @@ def main(argv=None) -> int:
         return cmd_worker(args.run)
     if args.command == "_head":
         return cmd_head(args.args_json)
+    if args.command == "rl":
+        from .rl.run import run_rl
+
+        run_rl(args)
+        return 0
     parser.print_help()
     return 0
 

yeto/rl/__init__.py

@@ -0,0 +1,2 @@
+from .envs.cybergym_env import CyberGymEnv
+from .algorithms.ppo import PPOTrainer