Initial commit: GRPO algorithm implementation

Author: SubaashNair

.github/workflows/publish.yml

@@ -0,0 +1,85 @@
+name: Publish Python distribution to PyPI
+
+on:
+  push:
+    tags:
+      - "v*.*.*"
+
+jobs:
+  build:
+    name: Build distribution 📦
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v4
+    - name: Set up Python
+      uses: actions/setup-python@v5
+      with:
+        python-version: "3.x"
+    - name: Install build tools
+      run: |
+        python -m pip install --upgrade pip
+        pip install build
+    - name: Build package
+      run: python -m build
+    - name: Store distribution packages
+      uses: actions/upload-artifact@v4
+      with:
+        name: python-package-distributions
+        path: dist/
+
+  publish-to-pypi:
+    name: Publish to PyPI
+    needs:
+    - build
+    runs-on: ubuntu-latest
+    environment:
+      name: pypi
+      url: https://pypi.org/p/optimrl
+    permissions:
+      id-token: write
+    
+    steps:
+    - name: Download built packages
+      uses: actions/download-artifact@v4
+      with:
+        name: python-package-distributions
+        path: dist/
+    - name: Publish to PyPI
+      uses: pypa/gh-action-pypi-publish@release/v1
+
+  github-release:
+    name: Create GitHub Release
+    needs:
+    - publish-to-pypi
+    runs-on: ubuntu-latest
+    permissions:
+      contents: write
+      id-token: write
+    
+    steps:
+    - name: Download packages
+      uses: actions/download-artifact@v4
+      with:
+        name: python-package-distributions
+        path: dist/
+    - name: Sign with Sigstore
+      uses: sigstore/[email protected]
+      with:
+        inputs: >-
+          ./dist/*.tar.gz
+          ./dist/*.whl
+    - name: Create Release
+      env:
+        GITHUB_TOKEN: ${{ github.token }}
+      run: >-
+        gh release create
+        "$GITHUB_REF_NAME"
+        --repo "$GITHUB_REPOSITORY"
+        --notes ""
+    - name: Upload assets
+      env:
+        GITHUB_TOKEN: ${{ github.token }}
+      run: >-
+        gh release upload
+        "$GITHUB_REF_NAME" dist/**
+        --repo "$GITHUB_REPOSITORY"

.github/workflows/tests.yml

@@ -0,0 +1,53 @@
+name: Tests
+
+on:
+  push:
+    branches: [ main ]
+  pull_request:
+    branches: [ main ]
+
+jobs:
+  test:
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        os: [ubuntu-latest, macos-latest, windows-latest]
+        python-version: ['3.8', '3.9', '3.10', '3.11', '3.12']
+
+    steps:
+    - uses: actions/checkout@v4
+    
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v5
+      with:
+        python-version: ${{ matrix.python-version }}
+    
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install -e '.[test]'
+    
+    - name: Run tests
+      run: |
+        pytest tests/ -v
+
+  lint:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v4
+    
+    - name: Set up Python
+      uses: actions/setup-python@v5
+      with:
+        python-version: '3.12'
+    
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install black isort flake8
+    
+    - name: Check code formatting
+      run: |
+        black --check optimrl tests
+        isort --check-only optimrl tests
+        flake8 optimrl tests

.gitignore

@@ -4,13 +4,19 @@ env/
 .env/
 .venv/
 ENV/
+optimrl/
+**/lib/
+**/bin/
+**/include/
+**/share/
+pyvenv.cfg
 
 # PyTorch model files
 *.pt
 *.pth
 *.ckpt
 
-# Python cache
+# Python cache 
 __pycache__/
 *.py[cod]
 *$py.class
@@ -19,16 +25,16 @@ __pycache__/
 dist/
 build/
 *.egg-info/
+*.pyc
+*.pyo
+*.so
 
 # IDE settings
 .vscode/
 .idea/
 *.swp
 *.swo
 
-# Jupyter Notebook
-.ipynb_checkpoints
-
 # Local development settings
 .env
 .env.local

OptimRL.code-workspace

@@ -0,0 +1,8 @@
+{
+	"folders": [
+		{
+			"path": "."
+		}
+	],
+	"settings": {}
+}

pyproject.toml

@@ -0,0 +1,7 @@
+[versioneer]
+VCS = git
+style = pep440
+versionfile_source = optimrl/_version.py
+versionfile_build = optimrl/_version.py
+tag_prefix = v
+parentdir_prefix = optimrl-

setup.cfg

@@ -0,0 +1,102 @@
+from setuptools import setup, find_packages, Extension
+from setuptools.command.build_ext import build_ext
+import os
+import platform
+import shutil
+import versioneer
+
+class CustomBuildExt(build_ext):
+    """Custom build command to handle library file copying and platform-specific compilation."""
+    
+    def build_extensions(self):
+        # Set platform-specific compiler flags
+        compiler_type = self.compiler.compiler_type
+        
+        for ext in self.extensions:
+            if compiler_type == 'unix':
+                # macOS and Linux specific flags
+                if platform.system() == 'Darwin':
+                    # macOS specific flags
+                    ext.extra_compile_args += ['-O3', '-fPIC']
+                    # Support both Intel and Apple Silicon
+                    ext.extra_compile_args += ['-arch', 'x86_64', '-arch', 'arm64']
+                    ext.extra_link_args += ['-arch', 'x86_64', '-arch', 'arm64']
+                else:
+                    # Linux specific flags
+                    ext.extra_compile_args += ['-O3', '-fPIC']
+            elif compiler_type == 'msvc':
+                # Windows specific flags
+                ext.extra_compile_args += ['/O2']
+
+        # Build the extensions
+        super().build_extensions()
+        
+        # After building, copy the library to the package directory
+        self.copy_extensions_to_package()
+    
+    def copy_extensions_to_package(self):
+        """Copy the built extension to the package directory."""
+        for ext in self.extensions:
+            # Get the full path of the built library
+            built_lib = self.get_ext_fullpath(ext.name)
+            
+            # Determine the destination directory within the package
+            dest_dir = os.path.join('optimrl', 'c_src')
+            os.makedirs(dest_dir, exist_ok=True)
+            
+            # Get the filename only
+            filename = os.path.basename(built_lib)
+            
+            # Create platform-specific library name
+            if platform.system() == 'Darwin':
+                lib_name = 'libgrpo.dylib'
+            elif platform.system() == 'Linux':
+                lib_name = 'libgrpo.so'
+            else:
+                lib_name = 'libgrpo.dll'
+            
+            # Copy the file to the package directory with the correct name
+            dest_path = os.path.join(dest_dir, lib_name)
+            shutil.copy2(built_lib, dest_path)
+            print(f"Copied {built_lib} to {dest_path}")
+
+# Define the extension module
+grpo_module = Extension(
+    'optimrl.c_src.libgrpo',
+    sources=['optimrl/c_src/grpo.c'],
+    include_dirs=['optimrl/c_src'],
+    libraries=['m'] if platform.system() != 'Windows' else [],
+    extra_compile_args=['-O3', '-fPIC'] if platform.system() != 'Windows' else ['/O2']
+)
+
+# Read the README file
+with open('README.md', 'r', encoding='utf-8') as f:
+    long_description = f.read()
+
+setup(
+    name="optimrl",
+    version=versioneer.get_version(),
+    cmdclass=versioneer.get_cmdclass(),
+    author="Subashanan Nair",
+    author_email="[email protected]",
+    description="Group Relative Policy Optimization for Efficient RL Training",
+    long_description=long_description,
+    long_description_content_type="text/markdown",
+    url="https://github.com/SubaashNair/OptimRL",
+    packages=find_packages(),
+    ext_modules=[grpo_module],
+    install_requires=[
+        "numpy>=1.20.0",
+        "torch>=1.8.0"
+    ],
+    extras_require={
+        'test': ['pytest>=6.0'],
+        'dev': ['pytest>=6.0', 'black', 'isort', 'flake8']
+    },
+    python_requires=">=3.8",
+    cmdclass={'build_ext': CustomBuildExt},
+    include_package_data=True,
+    package_data={
+        'optimrl': ['c_src/*.dylib', 'c_src/*.so', 'c_src/*.dll']
+    },
+)