create-rl-app 0.0.1__tar.gz

create_rl_app-0.0.1/PKG-INFO

@@ -0,0 +1,63 @@
+Metadata-Version: 2.3
+Name: create-rl-app
+Version: 0.0.1
+Summary: Add your description here
+Author: Koen Ponse
+Author-email: Koen Ponse <k.ponse@liacs.leidenuniv.nl>
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+
+# create-rl-app
+
+[![Python Version](https://img.shields.io/badge/python-3.11%2B-blue.svg)](https://python.org)
+[![Jaxnasium Version](https://badge.fury.io/py/jaxnasium.svg)](https://github.com/ponseko/jaxnasium)
+[![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
+
+A CLI application to bootstrap reinforcement learning applications within the Jaxnasium ecosystem. Quickly scaffold new RL projects for either developing environments with baseline algorithms or for altering existing baselines.
+
+## What it does
+
+`create-rl-app` is a command-line tool that helps you quickly set up new reinforcement learning projects using the Jaxnasium framework. It creates a well-structured project template with:
+
+- 🚀 **Quick Setup**: Get a new RL project running in seconds
+- 🏗️ **Helpful Templates**: Templates for environments and algorithms for you to start with.
+- ⚡ **Performance Optimized**: Sets you up with PureJaxRL compatible agents and environments for performance and GPU scalability.
+
+## Useage
+
+### uvx (Recommended)
+
+```bash
+uvx create-rl-app <project_name>
+cd <project_name>
+uv run train_example.py
+```
+
+### pipx
+
+```bash
+pipx run create-rl-app <project_name>
+cd <project_name>
+# Create a new environment (e.g. conda, venv, etc.)
+# source .../bin/activate
+python train_example.py
+```
+
+### Or Install Globally
+
+```bash
+uv tool install create-rl-app
+```
+
+```bash
+pip install create-rl-app
+```
+
+## Dependencies
+
+None.
+
+## Contributing
+
+Contributions are welcome! Please feel free to submit a Pull Request.
+

create_rl_app-0.0.1/README.md

@@ -0,0 +1,54 @@
+# create-rl-app
+
+[![Python Version](https://img.shields.io/badge/python-3.11%2B-blue.svg)](https://python.org)
+[![Jaxnasium Version](https://badge.fury.io/py/jaxnasium.svg)](https://github.com/ponseko/jaxnasium)
+[![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
+
+A CLI application to bootstrap reinforcement learning applications within the Jaxnasium ecosystem. Quickly scaffold new RL projects for either developing environments with baseline algorithms or for altering existing baselines.
+
+## What it does
+
+`create-rl-app` is a command-line tool that helps you quickly set up new reinforcement learning projects using the Jaxnasium framework. It creates a well-structured project template with:
+
+- 🚀 **Quick Setup**: Get a new RL project running in seconds
+- 🏗️ **Helpful Templates**: Templates for environments and algorithms for you to start with.
+- ⚡ **Performance Optimized**: Sets you up with PureJaxRL compatible agents and environments for performance and GPU scalability.
+
+## Useage
+
+### uvx (Recommended)
+
+```bash
+uvx create-rl-app <project_name>
+cd <project_name>
+uv run train_example.py
+```
+
+### pipx
+
+```bash
+pipx run create-rl-app <project_name>
+cd <project_name>
+# Create a new environment (e.g. conda, venv, etc.)
+# source .../bin/activate
+python train_example.py
+```
+
+### Or Install Globally
+
+```bash
+uv tool install create-rl-app
+```
+
+```bash
+pip install create-rl-app
+```
+
+## Dependencies
+
+None.
+
+## Contributing
+
+Contributions are welcome! Please feel free to submit a Pull Request.
+

create_rl_app-0.0.1/pyproject.toml

@@ -0,0 +1,17 @@
+[project]
+name = "create-rl-app"
+version = "0.0.1"
+description = "Add your description here"
+readme = "README.md"
+authors = [
+    { name = "Koen Ponse", email = "k.ponse@liacs.leidenuniv.nl" }
+]
+requires-python = ">=3.11"
+dependencies = []
+
+[project.scripts]
+create-rl-app = "create_rl_app.cli:main"
+
+[build-system]
+requires = ["uv_build>=0.8.4,<0.9.0"]
+build-backend = "uv_build"

create_rl_app-0.0.1/src/create_rl_app/_vendored/.gitkeep

@@ -0,0 +1,2 @@
+# This file ensures the vendored directory is tracked by git
+# even when empty. It will be removed when files are vendored.

create_rl_app-0.0.1/src/create_rl_app/cli.py

@@ -0,0 +1,317 @@
+import argparse
+import importlib.resources as pkg_resources
+import json
+import re
+import shutil
+import subprocess
+from pathlib import Path
+
+# ANSI color codes
+CYAN = "\033[96m"
+GREEN = "\033[92m"
+YELLOW = "\033[93m"
+BOLD = "\033[1m"
+RESET = "\033[0m"
+LIGHT_GRAY = "\033[37m"
+
+
+def get_vendored_jaxnasium_version():
+    """Get the vendored jaxnasium version from .vendor_info"""
+    try:
+        files = pkg_resources.files("create_rl_app._vendored")
+        vendor_info = files / "jaxnasium" / ".vendor_info"
+        with vendor_info.open("r") as f:
+            data = json.load(f)
+            version = data["version"]
+            return version.lstrip("v") if version.startswith("v") else version
+    except (FileNotFoundError, ImportError, KeyError):
+        raise RuntimeError("Vendored jaxnasium version not found")
+
+
+def colored_input(prompt, default=""):
+    """Display a colored input prompt with default value."""
+    default_display = f" ({LIGHT_GRAY}{default}{RESET})" if default else ""
+    user_input = input(
+        f"{CYAN}{BOLD}?{RESET} {prompt}{YELLOW}{default_display}{RESET}: "
+    ).strip()
+    return user_input if user_input else default
+
+
+def yes_no_prompt(question, default="y"):
+    """Ask a yes/no question with colored output."""
+    options = "(Y/n)" if default.lower() == "y" else "(y/N)"
+    while True:
+        response = colored_input(f"{question} {options}")
+        if not response:
+            return default.lower() == "y"
+        if response.lower() in ["y", "yes"]:
+            return True
+        elif response.lower() in ["n", "no"]:
+            return False
+        print(f"{YELLOW}Please answer with 'y' or 'n'.{RESET}")
+
+
+def update_pyproject_toml_file(project_path):
+    toml_path = project_path / "pyproject.toml"
+    if toml_path.exists():
+        with toml_path.open("r") as file:
+            content = file.read()
+
+        # Update requires-python version
+        content = re.sub(
+            r"^requires-python\s*=\s*.*$",
+            'requires-python = ">=3.11"',
+            content,
+            flags=re.MULTILINE,
+        )
+
+        # Update dependencies
+        def update_dependencies(match):
+            deps_str = match.group(1).strip()
+            if deps_str == "[]":
+                dependencies = []
+            else:
+                dependencies = re.findall(r'"([^"]+)"', deps_str)
+
+            # Add jaxnasium with the required version
+            dependencies.append(f"jaxnasium[algs]>={get_vendored_jaxnasium_version()}")
+            return f"dependencies = {dependencies}"
+
+        content = re.sub(
+            r"^dependencies\s*=\s*(\[.*?\])",
+            update_dependencies,
+            content,
+            flags=re.MULTILINE | re.DOTALL,
+        )
+
+        # UV-build expects a src folder by default. Remove this expectation here.
+        if re.search(r'build-backend\s*=\s*["\']uv_build["\']', content):
+            if not re.search(r"\[tool\.uv\.build-backend\]", content):
+                content += "\n\n[tool.uv.build-backend]\nmodule-root = ''\n"
+
+        with toml_path.open("w") as file:
+            file.write(content)
+
+
+def replace_init_file(src_path, import_env):
+    """Clear the project __init__.py file (or create one if it does not exist)
+    Fill it with "from projectname import ExampleEnv".
+    """
+    init_file_path = src_path / "__init__.py"
+    with init_file_path.open("w") as file:
+        file.write("# This file is auto-generated by jaxnasium\n")
+        if import_env:
+            file.write(f"from .{src_path.name}Env import ExampleEnv\n")
+
+
+def copy_env_template(target_path: Path):
+    target_path.parent.mkdir(parents=True, exist_ok=True)
+    with pkg_resources.path(
+        "create_rl_app.resources", "env_template.py"
+    ) as template_path:
+        shutil.copy(template_path, target_path)
+
+
+def copy_algorithms_source(target_path: Path):
+    target_path.mkdir(parents=True, exist_ok=True)
+    with pkg_resources.path(
+        "create_rl_app._vendored.jaxnasium", "algorithms"
+    ) as template_path:
+        # ignore the "utils" folder in "algorithms"
+        ignored = ["utils", "_algorithm.py"]
+        shutil.copytree(
+            template_path,
+            target_path,
+            dirs_exist_ok=True,
+            ignore=shutil.ignore_patterns(*ignored),
+        )
+
+
+def main():
+    # Parse the project name argument
+    parser = argparse.ArgumentParser(description="Initialize a new jaxnasium project.")
+    parser.add_argument("projectname", help="The path to the new project directory.")
+    parser.add_argument(
+        "-y",
+        "--yes",
+        action="store_true",
+        help="Automatically answer 'yes' to all prompts.",
+    )
+    parser.add_argument(
+        "--env-template",
+        type=str,
+        choices=["true", "false"],
+        help="Include environment template in the project (true/false).",
+    )
+    parser.add_argument(
+        "--algorithm-source",
+        type=str,
+        choices=["true", "false"],
+        help="Copy algorithm source code into the project instead of importing (true/false).",
+    )
+    parser.add_argument(
+        "--init-algorithms",
+        action="store_true",
+        help="Skip project creation and only copy algorithm source code at the <projectname> location",
+    )
+    parser.add_argument(
+        "--init-env",
+        action="store_true",
+        help="Skip project creation and only copy environment source code at the <projectname> location",
+    )
+    args = parser.parse_args()
+    projectname = args.projectname
+
+    if args.init_algorithms:
+        copy_algorithms_source(Path(projectname))
+
+    if args.init_env:
+        copy_env_template(Path(projectname))
+
+    if args.init_algorithms or args.init_env:
+        return
+
+    print(f"{CYAN}{BOLD}")
+    print("  ██████╗██████╗ ███████╗ █████╗ ████████╗███████╗")
+    print(" ██╔════╝██╔══██╗██╔════╝██╔══██╗╚══██╔══╝██╔════╝")
+    print(" ██║     ██████╔╝█████╗  ███████║   ██║   █████╗")
+    print(" ██║     ██╔══██╗██╔══╝  ██╔══██║   ██║   ██╔══╝")
+    print(" ╚██████╗██║  ██║███████╗██║  ██║   ██║   ███████╗")
+    print("  ╚═════╝╚═╝  ╚═╝╚══════╝╚═╝  ╚═╝   ╚═╝   ╚══════╝")
+    print("")
+    print(" ██████╗ ██╗               █████╗ ██████╗ ██████╗")
+    print(" ██╔══██╗██║              ██╔══██╗██╔══██╗██╔══██╗")
+    print(" ██████╔╝██║              ███████║██████╔╝██████╔╝")
+    print(" ██╔══██╗██║              ██╔══██║██╔═══╝ ██╔═══╝")
+    print(" ██║  ██║███████╗         ██║  ██║██║     ██║")
+    print(" ╚═╝  ╚═╝╚══════╝         ╚═╝  ╚═╝╚═╝     ╚═╝")
+    print(f"{RESET}")
+
+    print(
+        f"{LIGHT_GRAY}Setting up a new Jaxnasium project (v{get_vendored_jaxnasium_version()}){RESET}"
+    )
+    print(f"{LIGHT_GRAY}{'─' * 60}{RESET}\n")
+
+    if any(char.isupper() for char in projectname):
+        print(f"{YELLOW} Project name has been altered to lowercase{RESET}")
+    projectname = projectname.lower()
+    print(f"{LIGHT_GRAY}Project name: {projectname}{RESET}")
+
+    ########## Questions ##########
+
+    # Use CLI arguments if provided, otherwise prompt the user
+    if args.env_template is not None:
+        build_environment_template = args.env_template == "true"
+    else:
+        build_environment_template = args.yes or yes_no_prompt(
+            "Would you like to include a environment template?", default="y"
+        )
+
+    if args.algorithm_source is not None:
+        include_algorithm_source = args.algorithm_source == "true"
+    else:
+        include_algorithm_source = args.yes or yes_no_prompt(
+            "Instead of importing, would you like to copy the algorithm source code into your project?",
+            default="n",
+        )
+
+    # Display summary of choices
+    print(f"\n{BOLD}📋 Project configuration summary:{RESET}")
+    print(f"  • Project name: {projectname}")
+    print(
+        f"  • Include environment template: {'Yes' if build_environment_template else 'No'}"
+    )
+    print(
+        f"  • Copy algorithm source code: {'Yes' if include_algorithm_source else 'No'}"
+    )
+
+    # Confirm setup
+    if not args.yes:
+        if not yes_no_prompt("\nDo you want to proceed with this configuration?"):
+            print("Setup cancelled.")
+            return
+
+    #### SETUP PROJECT ####
+
+    # Determine the command to run
+    command = (
+        ["uv", "init", "--package", projectname]
+        if shutil.which("uv")
+        else ["pipx", "run", "uv", "init", "--package", projectname]
+    )
+
+    # Run UV init command
+    try:
+        subprocess.run(command, check=True)
+    except subprocess.CalledProcessError as e:
+        raise RuntimeError(
+            f"Failed to initialize the project: {e}. \n Most likely neither uv nor pipx are installed."
+        )
+
+    # Set project paths
+    project_path = Path(projectname).resolve()
+    src_path = project_path / "src" / project_path.name
+
+    # UV defaults to src structure. Create a flat structure instead:
+    if src_path.exists():
+        shutil.move(src_path, project_path / project_path.name)
+        if (project_path / "src").exists():
+            shutil.rmtree(project_path / "src")
+    src_path = project_path / project_path.name
+    assert src_path.exists(), "Src path does not exist"
+
+    # Update the dependencies in pyproject.toml and set the minimum python version
+    update_pyproject_toml_file(project_path)
+
+    # Update default __init__.py file generated by uv
+    replace_init_file(src_path, build_environment_template)
+
+    # Copy over the required files.
+    if build_environment_template:
+        file_path = src_path / f"{project_path.name}Env.py"
+        copy_env_template(file_path)
+
+    if include_algorithm_source:
+        file_path = src_path / "algorithms"
+        copy_algorithms_source(file_path)
+
+    with pkg_resources.path(
+        "create_rl_app.resources", "train_template.py"
+    ) as template_path:
+        shutil.copy(template_path, project_path / "train_example.py")
+
+        # Read template content and apply regex replacements
+        with open(template_path, "r") as file:
+            content = file.read()
+
+        # Replace import statements and environment creation based on configuration
+        if build_environment_template:
+            # Replace the jym.make call to use ExampleEnv
+            content = re.sub(
+                r'jym\.make\("CartPole-v1"\)',
+                'jym.make("ExampleEnv")',
+                content,
+                flags=re.MULTILINE,
+            )
+            # Add import statement of example environment
+            import_line = (
+                f"from {project_path.name}.{project_path.name}Env import ExampleEnv\n"
+            )
+            content = import_line + content
+
+        # Replace algorithm import if copying source code
+        if include_algorithm_source:
+            content = re.sub(
+                r"^from jaxnasium\.algorithms.*$",
+                f"from {project_path.name}.algorithms import PPO",
+                content,
+                flags=re.MULTILINE,
+            )
+
+        # Write the modified content
+        with open(project_path / "train_example.py", "w") as file:
+            file.write(content)
+
+
+if __name__ == "__main__":
+    main()

create_rl_app-0.0.1/src/create_rl_app/resources/env_template.py

@@ -0,0 +1,89 @@
+from typing import Tuple
+
+import equinox as eqx
+import jax.numpy as jnp
+import jaxnasium as jym
+from jaxtyping import Array, Float, PRNGKeyArray
+
+
+class EnvState(eqx.Module):
+    x: int
+    y: int
+    time: int = 0
+
+    @property
+    def location(self):
+        return (self.x, self.y)
+
+
+@jym.registry.register("ExampleEnv")
+class ExampleEnv(jym.Environment):
+    max_episode_steps: int = 100
+
+    def step_env(
+        self, key: PRNGKeyArray, state: EnvState, action: int
+    ) -> Tuple[jym.TimeStep, EnvState]:
+        """
+        Update the environment state based on the action taken.
+        """
+        # action 0 -> move up
+        # action 1 -> move right
+        # action 2 -> move down
+        # action 3 -> move left
+        new_x = state.x + (action == 0) - (action == 2)
+        new_y = state.y + (action == 1) - (action == 3)
+
+        state = EnvState(x=new_x, y=new_y, time=state.time + 1)
+
+        timestep = jym.TimeStep(
+            observation=self.get_observation(state),
+            reward=self.get_reward(),
+            terminated=self.get_terminated(state),
+            truncated=state.time >= self.max_episode_steps,
+            info={},
+        )
+        return timestep, state
+
+    def reset_env(self, key: PRNGKeyArray) -> Tuple[Array, EnvState]:
+        """
+        Reset the environment to its initial state.
+        """
+        state = EnvState(x=5, y=5)  # Start in the center
+        observation = self.get_observation(state)
+        return observation, state
+
+    def get_observation(self, state: EnvState) -> Array:
+        """
+        Get the observation from the environment state.
+        """
+        return jnp.array(state.location)
+
+    def get_reward(self) -> float:
+        """
+        Get the reward from the environment state.
+        """
+        # Example reward function: 1 for each step taken
+        return 1.0
+
+    def get_terminated(self, state: EnvState) -> Float[Array, "1"]:
+        """
+        Check if the episode has terminated.
+        """
+        # Example termination condition: if the agent moves out of bounds
+        out_of_bounds_x = jnp.logical_or(state.x < 0, state.x >= 10)
+        out_of_bounds_y = jnp.logical_or(state.y < 0, state.y >= 10)
+        return jnp.logical_or(out_of_bounds_x, out_of_bounds_y)
+
+    @property
+    def observation_space(self) -> jym.Space:
+        """
+        Define the observation space of the environment.
+        """
+        return jym.Box(low=0, high=10, shape=(2,))
+
+    @property
+    def action_space(self) -> jym.Space:
+        """
+        Define the action space of the environment.
+        """
+        return jym.Discrete(n=4)

create_rl_app-0.0.1/src/create_rl_app/resources/train_template.py

@@ -0,0 +1,43 @@
+import jax
+import jaxnasium as jym
+from jaxnasium.algorithms import PPO
+from jaxtyping import PRNGKeyArray
+
+
+def do_random_evaluation(
+    key: PRNGKeyArray, env: jym.Environment, num_repitions: int = 10
+):
+    """Perform some random steps to set a baseline for the environment."""
+    rewards = 0.0
+    for _ in range(num_repitions):
+        obs, env_state = env.reset(key)
+        while True:
+            key, key = jax.random.split(key)
+            action = env.action_space.sample(key)
+            (obs, reward, terminated, truncated, info), env_state = env.step(
+                key, env_state, action
+            )
+            rewards += reward
+            if terminated or truncated:
+                break
+    return rewards / num_repitions
+
+
+if __name__ == "__main__":
+    env = jym.make("CartPole-v1")
+    env = jym.LogWrapper(env)
+    rng = jax.random.PRNGKey(0)
+
+    random_rewards = do_random_evaluation(rng, env)
+    print(f"Random Agent average reward: {random_rewards}")
+
+    # RL Training with PPO
+    agent = PPO(
+        total_timesteps=50000,
+        num_steps=64,
+        learning_rate=2.5e-3,
+        ent_coef=0.0,
+        num_epochs=1,
+    )
+    agent = agent.train(rng, env)
+    print(f"Agent average reward: {agent.evaluate(rng, env, num_eval_episodes=5)}")