[BUG] Stable-Baselines3 integration not working

[eltonjohnfanboy]

[BUG] Stable-Baselines3 integration not working

[mihran113]

Hey @eltonjohnfanboy! Sorry for late response and thanks for opening the issue. There's an issue in our callback, I'll make sure to include the fix in the upcoming release. In the meantime you can use the following script as a workaround to track the metrics:

from typing import Any, Dict, Tuple, Union
import os
import gymnasium
from stable_baselines3 import PPO
from stable_baselines3.common.vec_env import DummyVecEnv
from stable_baselines3.common.evaluation import evaluate_policy
from stable_baselines3.common.logger import HumanOutputFormat, KVWriter, Logger

from aim import Run
import numpy as np


class AimOutputFormat(KVWriter):
    """
    Track key/value pairs into Aim run.
    """

    def __init__(
        self,
        aim_run
    ):
        self.aim_run = aim_run

    def write(
        self,
        key_values: Dict[str, Any],
        key_excluded: Dict[str, Union[str, Tuple[str, ...]]],
        step: int = 0,
    ) -> None:
        for (key, value), (_, excluded) in zip(
            sorted(key_values.items()), sorted(key_excluded.items())
        ):
            if excluded is not None and 'aim' in excluded:
                continue

            if isinstance(value, np.ScalarType):
                if not isinstance(value, str):
                    tag, key = key.split('/')
                    if tag in ['train', 'valid']:
                        context = {'subset': tag}
                    else:
                        context = {'tag': tag}

                    self.aim_run.track(value, key, step=step, context=context)


run = Run()
loggers = Logger(
    folder=None,
    output_formats=[AimOutputFormat(run)],
)

environment_name = 'CartPole-v1'
env = gymnasium.make(environment_name)

env = DummyVecEnv([lambda: env])
model = PPO('MlpPolicy', env, verbose = 1)
model.set_logger(loggers)
model.learn(total_timesteps=10_000)

[eltonjohnfanboy]

Hi @mihran113, no worries at all. Thanks a lot for your work and the update! Looking forward to the upcoming release and will be using the script in the meantime. :))

[mihran113]

Hey @eltonjohnfanboy! The new version of aim has been shipped (3.23.0) which includes the fix for this issue. Please let me know if everything works as expected so I can close the issue.

[eltonjohnfanboy]

Hi @mihran113, sorry for the delayed response. Thanks for the update; it works correctly now! :))