[Tune] [PBT] [Doc] Add example PBT notebook

[justinvyu]

Summary

Example notebook

See here for a Colab version of the notebook! Some images are missing that would show up in the Ray docs once merged.

The purpose of the example is to help new users understand what PBT is doing under the hood, and provide an example of using PBT with a function trainable. The notebook gives recommendations on how to set PBT-specific parameters such as perturbation_interval and how to perform checkpointing alongside PBT.

The example notebook replicates an experiment found in the original PBT paper. The toy example in the paper is a good way of verifying that PBT behavior is correct (trials are being correctly exploited and the correct checkpoints are being used).

Notebook plots

Figure in paper

PBT logging enhancements

This PR also includes some quality of life improvements to PBT logging that makes it easier to understand what is happening.

Example:

hyperparam_mutations = {
    "a": tune.uniform(0, 1),
    "b": list(range(5)),
    "c": {
        "d": tune.uniform(2, 3),
        "e": {"f": [-1, 0, 1]},
    },
}

The above hyperparameter mutations config results in the following log after the PBT explore step:

[PBT] [Explore] Perturbed the hyperparameter config of trial trial_1:
a : 0.5 --- (* 1.2) --> 0.6
b : 2 --- (shift right) --> 3
c :
    d : 2.5 --- (* 1.2) --> 3.0
    e :
        f : 0 --- (shift right) --> 1

Some docs restructuring

I also took the chance to nest some Tune examples in sub-sections, since there were a lot of frameworks specific examples that were overloading the table of contents.

Related PRs

This test helped uncover some bugs fixed in the following PRs:

• [Tune] [PBT] Maintain consistent Trial/TrialRunner state when pausing and resuming trial #28511
• [Tune] [PBT] Ensure that exploited in-memory checkpoint is used by trial #28509

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[xwjiang2010]

Thanks so much!
This looks great! Only some nits..

[xwjiang2010]

update?

[justinvyu]

What should be updated? The return type should be updated already.

[xwjiang2010]

there is also operations returned.

[xwjiang2010]

mention this is to illustrate synchronous pbt?

[krfricke]

Amazing tutorial and changes! Just two minor nits.

cc @maxpumperla for docs approval (incl. structure)

[krfricke]

Can we add a comment here that explains what we're doing?

[krfricke]

Generally, let's add a few comments for this whole exploration block

[krfricke]

Should this be "shift left (noop)" or similar if old_idx == new_idx? E.g. when we select shift = -1 when we're already at the first item

[justinvyu]

Right, this should include some indicator that we've hit the end of the list.

A few thoughts on this actually:

• We might want to just wrap around from the end if trying to shift left at the first item. The fact that we don't wrap implies some kind of ordering to the items, but there is no guarantee there. Example:
  • "a": [100, 1, 50, 75]
  • 75 <- (shift left & wrap around) -- 100 -- (shift right) -> 1
  • Seems arbitrary why 100 can only be perturbed to the right.
• Should a noop be an option here rather than always shifting left and right? SHERPA's implementation includes a "no shift" option.

[krfricke]

Hm what does the original paper say?

Everything else seems to have an option not to shift, so I actually believe just going with [-1, 0, 1] should be fine.

For wrapping around, basically a matter of preference. Best would be to distinguish between ordinal and categorical, but since we don't have this, I'd say let's assume ordinal and leave the logic as is

[justinvyu]

1. The original paper doesn't really mention specifics on perturb, so I think adding the no-shift makes sense to align with other implementations.

2. I see, can leave this for a PR in the future.