[rllib, train] Add support for nested metrics in Result.get_best_checkpoint

[pseudo-rnd-thoughts]

Description

RLlib uses nested metric structure (like "{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}") which Result.get_best_checkpoint doesn't support.
Following ResultGrid.get_best_result() to use unflattened_lookup, I've added that to get_best_checkpoint along with testing for nested structures (and its backward compatibility)

Reproduction script

from ray import tune
from ray.rllib.algorithms.ppo import PPOConfig
from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig
from ray.rllib.utils.metrics import (
    ENV_RUNNER_RESULTS,
    EPISODE_RETURN_MEAN,
    NUM_ENV_STEPS_SAMPLED_LIFETIME,
)
from ray.tune.result import TRAINING_ITERATION

config = (
    PPOConfig()
    .environment("CartPole-v1")
    .training(num_epochs=6)
)

tuner = tune.Tuner(
    "PPO",
    param_space=config.to_dict(),
    run_config=tune.RunConfig(
        "PPO_Reproduce",
        checkpoint_config=tune.CheckpointConfig(
            num_to_keep=10,
            checkpoint_score_attribute=f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}",
            checkpoint_at_end=True,
            checkpoint_frequency=5,
        ),
        stop={
            f"{ENV_RUNNER_RESULTS}/{NUM_ENV_STEPS_SAMPLED_LIFETIME}": 3e5,
            f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": 450,
            TRAINING_ITERATION: 100,
        },
    ),
)

results = tuner.fit()
best_result = results.get_best_result()
ckpt = best_result.get_best_checkpoint(f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}", "max")
print(ckpt.path)

Related issues

#57533

[gemini-code-assist]

Code Review

This pull request adds support for nested metrics in Result.get_best_checkpoint by using unflattened_lookup. The changes are correct and are accompanied by good tests covering nested metrics, different modes, and backward compatibility. I've suggested one improvement to the error message when an invalid metric is provided, to make it more helpful for users of nested metrics.

[matthewdeng]

@justinvyu do we want to support this for Train V2 as well, or should we diverge for Tune?

[justinvyu]

I'm ok to support this in Train. This is only needed if users self-report nested dicts.

[justinvyu]

I'm ok to support this in Train. This is only needed if users self-report nested dicts.

[justinvyu]

Thanks!