Tune/WandB - Group different samples under job type

Ray Libraries (Data, Train, Tune, Serve)
1

Hey all,
I am logging my RLLIB experiments via the Weight and Biases integrations which has been working great so far. However, I would like to be able to log all experiments with the same config (so just different samples, via num_samples in the tune_config) under the same job type. Is this possible? See a minimal example below.

import ray
from ray import air, tune
from ray.air.integrations.wandb import WandbLoggerCallback
from ray.rllib.algorithms.ppo import PPOConfig
from ray.rllib.agents.ppo import PPOTrainer
import pyglet
pyglet.options["headless"] = True
import gymnasium as gym

if __name__ == "__main__":
    def trial_name_string(trial) -> str:
        trial_name = f"{trial.trial_id}"
        return trial_name

    ray.init(local_mode=True, num_cpus=10, num_gpus=0)
    stop = {
        "training_iteration": 2
    }
    env_name = "CartPole-v1"
    env = gym.make(env_name)
    config = {
        "num_workers": 1,
        "env": "CartPole-v1",
        'gamma': tune.grid_search([0.99, 0.95, 0.9, 0.99]),
        "framework": "torch",
        "log_level": "CRITICAL",
        "num_gpus": 0,
        "disable_env_checking": True
    }

    DEFAULT_CONFIG = PPOConfig().to_dict()
    config = {**DEFAULT_CONFIG, **config}

    tuner = ray.tune.Tuner(
        PPOTrainer,
        param_space=config,
        run_config=air.RunConfig(
            stop=stop,
            local_dir="ray_results",
            verbose=1,
            callbacks=[
                WandbLoggerCallback(
                    project=f"[project]",
                    entity="[name]",
                    group=f"[setting]",
                    job_type="[grouped by same setting for different samples]", #<- func. similar to trial_name_string?
                    log_config=True,
                    save_code=True,
                    monitor_gym=True
                )
            ],
            failure_config=air.FailureConfig(
                max_failures=2
            )
        ),
        tune_config=tune.TuneConfig(
            num_samples=10,
            trial_name_creator=trial_name_string,
            trial_dirname_creator=trial_name_string,
        ),
    )
    results = tuner.fit()
    ray.shutdown()
2

Hi @Thomas_Langerak1,

Thanks for the detailed post! This is not currently supported for the built-in WandbLoggerCallback, but it is a good feature request that we can support in the future. I’ve created an issue to track this here: [AIR][Tune] Add an option in `WandbLoggerCallback` to group wandb runs by config · Issue #33084 · ray-project/ray · GitHub

For now, you can use the setup_wandb method of using wandb with Tune instead. When setting up your PPOTrainer (using some RLlib callback), you can call setup_wandb, and then log things with a group set to something common for each config. See this guide for an example of setup_wandb: Using Weights & Biases with Tune — Ray 2.3.0

Here’s the API reference: External library integrations for Ray Tune — Ray 2.3.0

Let me know if this works out for you!