{
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.10.0"
  },
  "colab": {
   "provenance": [],
   "toc_visible": true
  }
 },
 "cells": [
  {
   "cell_type": "markdown",
   "id": "title",
   "metadata": {},
   "source": [
    "# DeepGym Quickstart\n",
    "\n",
    "**DeepGym** — RL training environments with verifiable rewards for coding agents.\n",
    "\n",
    "Give it model-generated code. It runs the code, verifies it, returns a score.\n",
    "That score plugs into TRL, verl, OpenRLHF, or any RL training loop.\n",
    "\n",
    "[![GitHub](https://img.shields.io/badge/GitHub-DeepGym-black)](https://github.com/DeepGym/deepgym)\n",
    "[![PyPI](https://img.shields.io/pypi/v/deepgym)](https://pypi.org/project/deepgym/)\n",
    "\n",
    "This notebook covers:\n",
    "1. Score a single solution\n",
    "2. Batch scoring (GRPO-style)\n",
    "3. Per-test-case reward traces\n",
    "4. Drop into TRL GRPOTrainer\n",
    "5. Write your own verifier\n",
    "6. Browse built-in environments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "install",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install DeepGym\n",
    "!pip install deepgym -q"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "section-1",
   "metadata": {},
   "source": [
    "## 1. Score a single solution\n",
    "\n",
    "Load a built-in environment, submit a solution, get a score."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "single-score",
   "metadata": {},
   "outputs": [],
   "source": [
    "from deepgym import DeepGym, load_environment\n",
    "\n",
    "dg = DeepGym(mode='local')          # no Daytona needed for local dev\n",
    "env = load_environment('coin_change')\n",
    "\n",
    "correct_solution = '''\n",
    "def coin_change(coins, amount):\n",
    "    dp = [float('inf')] * (amount + 1)\n",
    "    dp[0] = 0\n",
    "    for coin in coins:\n",
    "        for x in range(coin, amount + 1):\n",
    "            dp[x] = min(dp[x], dp[x - coin] + 1)\n",
    "    return dp[amount] if dp[amount] != float('inf') else -1\n",
    "'''\n",
    "\n",
    "result = dg.run(env, model_output=correct_solution)\n",
    "print(f'Score:  {result.score}')\n",
    "print(f'Passed: {result.passed}')\n",
    "print(f'Time:   {result.execution_time_ms:.0f}ms')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "wrong-solution",
   "metadata": {},
   "outputs": [],
   "source": [
    "# A wrong solution scores 0.0\n",
    "wrong_solution = '''\n",
    "def coin_change(coins, amount):\n",
    "    return -1  # always wrong\n",
    "'''\n",
    "\n",
    "result = dg.run(env, model_output=wrong_solution)\n",
    "print(f'Score:  {result.score}')   # 0.0\n",
    "print(f'Passed: {result.passed}')  # False"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "section-2",
   "metadata": {},
   "source": [
    "## 2. Batch scoring for GRPO\n",
    "\n",
    "GRPO needs to score N completions for the same prompt and compute advantages.\n",
    "DeepGym runs them in parallel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "batch-scoring",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simulate 8 model completions (mix of correct and wrong)\n",
    "solutions = [\n",
    "    # Correct DP solution\n",
    "    '''\n",
    "def coin_change(coins, amount):\n",
    "    dp = [float('inf')] * (amount + 1)\n",
    "    dp[0] = 0\n",
    "    for coin in coins:\n",
    "        for x in range(coin, amount + 1):\n",
    "            dp[x] = min(dp[x], dp[x - coin] + 1)\n",
    "    return dp[amount] if dp[amount] != float('inf') else -1\n",
    "''',\n",
    "    # Greedy (wrong for some inputs)\n",
    "    '''\n",
    "def coin_change(coins, amount):\n",
    "    coins.sort(reverse=True)\n",
    "    count = 0\n",
    "    for coin in coins:\n",
    "        count += amount // coin\n",
    "        amount %= coin\n",
    "    return count if amount == 0 else -1\n",
    "''',\n",
    "    # Always returns -1\n",
    "    'def coin_change(coins, amount): return -1',\n",
    "    # Correct BFS solution\n",
    "    '''\n",
    "from collections import deque\n",
    "def coin_change(coins, amount):\n",
    "    if amount == 0: return 0\n",
    "    visited = {0}\n",
    "    q = deque([(0, 0)])\n",
    "    while q:\n",
    "        curr, steps = q.popleft()\n",
    "        for coin in coins:\n",
    "            nxt = curr + coin\n",
    "            if nxt == amount: return steps + 1\n",
    "            if nxt < amount and nxt not in visited:\n",
    "                visited.add(nxt)\n",
    "                q.append((nxt, steps + 1))\n",
    "    return -1\n",
    "''',\n",
    "] * 2  # 8 total\n",
    "\n",
    "batch = dg.run_batch(env, solutions, max_parallel=8)\n",
    "scores = [r.score for r in batch.results]\n",
    "\n",
    "print(f'Scores: {[f\"{s:.2f}\" for s in scores]}')\n",
    "print(f'Passed: {batch.passed}/{batch.total}')\n",
    "print(f'Avg:    {batch.avg_score:.3f}')\n",
    "\n",
    "# GRPO advantage computation\n",
    "mean_score = sum(scores) / len(scores)\n",
    "std_score = (sum((s - mean_score) ** 2 for s in scores) / len(scores)) ** 0.5\n",
    "advantages = [(s - mean_score) / (std_score + 1e-8) for s in scores]\n",
    "print(f'Advantages: {[f\"{a:.2f}\" for a in advantages]}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "section-3",
   "metadata": {},
   "source": [
    "## 3. Per-test-case reward traces\n",
    "\n",
    "DeepGym returns **which specific tests passed and failed** — not just an overall score.\n",
    "This gives GRPO a denser training signal."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "per-test-traces",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = dg.run(env, model_output=correct_solution)\n",
    "\n",
    "if result.cases:\n",
    "    print(f'Total test cases: {len(result.cases)}')\n",
    "    print()\n",
    "    for case in result.cases[:5]:  # show first 5\n",
    "        status = 'PASS' if case.passed else 'FAIL'\n",
    "        print(f'  [{status}] {case.id}: {case.input_summary}')\n",
    "        if not case.passed and case.error:\n",
    "            print(f'         Error: {case.error}')\n",
    "else:\n",
    "    print('No per-test traces (verifier does not emit cases)')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "per-test-reward-shaping",
   "metadata": {},
   "outputs": [],
   "source": [
    "from deepgym import RewardFunction\n",
    "\n",
    "reward_fn = RewardFunction(env=env, max_parallel=8)\n",
    "\n",
    "# Per-test-case breakdown for each completion\n",
    "per_test = reward_fn.per_test_rewards(solutions[:4])\n",
    "for i, breakdown in enumerate(per_test):\n",
    "    print(f'Solution {i}: overall={breakdown[\"overall\"]:.2f}, '\n",
    "          f'tests={sum(1 for k,v in breakdown.items() if k != \"overall\" and v == 1.0)}/'\n",
    "          f'{sum(1 for k in breakdown if k != \"overall\")} passed')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "section-4",
   "metadata": {},
   "source": [
    "## 4. Drop into TRL GRPOTrainer\n",
    "\n",
    "One-liner integration with Hugging Face TRL."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "trl-integration",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install TRL if not present\n",
    "# !pip install trl -q\n",
    "\n",
    "from deepgym.integrations.trl import make_trl_reward_fn\n",
    "\n",
    "env = load_environment('two_sum')\n",
    "reward_fn = make_trl_reward_fn(env)  # returns Callable[[list[str]], list[float]]\n",
    "\n",
    "# Test it directly\n",
    "test_completions = [\n",
    "    'def two_sum(nums, target):\\n    seen = {}\\n    for i, n in enumerate(nums):\\n        if target - n in seen:\\n            return [seen[target - n], i]\\n        seen[n] = i',\n",
    "    'def two_sum(nums, target): return []',  # wrong\n",
    "]\n",
    "scores = reward_fn(test_completions)\n",
    "print(f'TRL reward scores: {scores}')\n",
    "\n",
    "# In a real training loop:\n",
    "# from trl import GRPOTrainer\n",
    "# trainer = GRPOTrainer(\n",
    "#     model='Qwen/Qwen2-0.5B-Instruct',\n",
    "#     reward_funcs=[reward_fn],\n",
    "#     train_dataset=dataset,\n",
    "# )\n",
    "# trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "section-5",
   "metadata": {},
   "source": [
    "## 5. Write your own verifier\n",
    "\n",
    "The verifier is just Python code that returns a float 0.0-1.0 (or bool or dict).\n",
    "DeepGym wraps it into the JSON protocol automatically."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "custom-verifier",
   "metadata": {},
   "outputs": [],
   "source": [
    "from deepgym import DeepGym, Environment\n",
    "\n",
    "# Verifier body: receives solution_path and optional test_cases_path\n",
    "# Return float, bool, or dict with {score, passed, details, cases}\n",
    "verifier_code = '''\n",
    "import importlib.util, sys\n",
    "spec = importlib.util.spec_from_file_location(\"sol\", solution_path)\n",
    "mod = importlib.util.module_from_spec(spec)\n",
    "spec.loader.exec_module(mod)\n",
    "\n",
    "cases = [\n",
    "    (2, 3, 5),\n",
    "    (0, 0, 0),\n",
    "    (-1, 1, 0),\n",
    "    (100, 200, 300),\n",
    "]\n",
    "\n",
    "passed = 0\n",
    "case_results = []\n",
    "for a, b, expected in cases:\n",
    "    try:\n",
    "        got = mod.add(a, b)\n",
    "        ok = got == expected\n",
    "    except Exception as e:\n",
    "        ok = False\n",
    "    if ok:\n",
    "        passed += 1\n",
    "    case_results.append({\n",
    "        \"id\": f\"test_{a}+{b}\",\n",
    "        \"passed\": ok,\n",
    "        \"score\": 1.0 if ok else 0.0,\n",
    "        \"input_summary\": f\"add({a}, {b})\",\n",
    "        \"expected_summary\": str(expected),\n",
    "    })\n",
    "\n",
    "score = passed / len(cases)\n",
    "return {\"score\": score, \"passed\": score == 1.0, \"details\": f\"{passed}/{len(cases)} passed\", \"cases\": case_results}\n",
    "'''\n",
    "\n",
    "env = Environment(\n",
    "    task='Write a function `add(a, b)` that returns the sum of two numbers.',\n",
    "    verifier_code=verifier_code,\n",
    ")\n",
    "\n",
    "dg = DeepGym(mode='local')\n",
    "result = dg.run(env, model_output='def add(a, b):\\n    return a + b')\n",
    "print(f'Score:  {result.score}')\n",
    "print(f'Passed: {result.passed}')\n",
    "for c in (result.cases or []):\n",
    "    print(f'  {c.id}: {\"PASS\" if c.passed else \"FAIL\"}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "section-6",
   "metadata": {},
   "source": [
    "## 6. Browse built-in environments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "list-envs",
   "metadata": {},
   "outputs": [],
   "source": [
    "from deepgym import list_environments\n",
    "\n",
    "envs = list_environments()\n",
    "print(f'{len(envs)} built-in environments:')\n",
    "for name in sorted(envs):\n",
    "    print(f'  {name}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "benchmark-envs",
   "metadata": {},
   "outputs": [],
   "source": [
    "# After importing benchmarks (run these scripts first):\n",
    "# !python scripts/import_humaneval.py      # 164 HumanEval problems\n",
    "# !python scripts/import_mbpp.py           # 500 MBPP problems\n",
    "# !python scripts/import_bigcodebench.py   # 1,140 BigCodeBench problems\n",
    "\n",
    "# Then:\n",
    "# env = load_environment('HumanEval/0')   # first HumanEval problem\n",
    "# env = load_environment('MBPP/1')        # first MBPP problem\n",
    "print('Run the import scripts above to get 2,350+ benchmark environments.')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "section-gymnasium",
   "metadata": {},
   "source": [
    "## 7. Gymnasium-style API\n",
    "\n",
    "If you prefer a reset/step loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "gymnasium",
   "metadata": {},
   "outputs": [],
   "source": [
    "from deepgym import DeepGymEnv, load_environment\n",
    "\n",
    "env = load_environment('reverse_string')\n",
    "gym_env = DeepGymEnv(env=env, mode='local')\n",
    "\n",
    "obs = gym_env.reset()\n",
    "print('Task:', obs.task[:80], '...')\n",
    "\n",
    "solution = '''\n",
    "def reverse_string(s):\n",
    "    return s[::-1]\n",
    "'''\n",
    "\n",
    "obs, reward, terminated, info = gym_env.step(solution)\n",
    "print(f'Reward: {reward:.2f}')\n",
    "print(f'Done:   {terminated}')"
   ]
  }
 ]
}
