stable_baselines3/train.py

#!/usr/bin/env python3

import mlflow
import dagshub
import gym
from gym import spaces
import numpy as np
from stable_baselines3 import PPO
import os

class PPOModelWrapper(mlflow.pyfunc.PythonModel):
    def load_context(self, context):
        self.model = PPO.load(context.artifacts["path"])

    def predict(self, context, model_input):
        action, states = self.model.predict(model_input)
        return {"action": action, "states": states}

# Create a simple custom gym environment
class SimpleEnv(gym.Env):
    def __init__(self):
        super(SimpleEnv, self).__init__()
        self.action_space = spaces.Discrete(3)
        self.observation_space = spaces.Box(low=-np.inf, high=np.inf, shape=(4,))

    def step(self, action):
        return np.array([0, 0, 0, 0]), 0, False, {}

    def reset(self):
        return np.array([0, 0, 0, 0])



if __name__ == '__main__':
    dagshub.init('stable_baselines3', 'jinensetpal') # repo name

    # Create and train the model
    env = SimpleEnv()
    model = PPO("MlpPolicy", env, verbose=1)
    model.learn(total_timesteps=1000)
    model.save('ppo')
    
    # Save the model using MLflow
    mlflow.pyfunc.log_model('model',
                            python_model=PPOModelWrapper(),
                            artifacts={'path': 'ppo.zip'})
    
    # Load the model from MLflow using the captured run_id
    run_id = mlflow.active_run().info.run_id
    loaded_model = mlflow.pyfunc.load_model(f"runs:/{run_id}/model")