[RLlib] Docs update: New Algorithm checkpoints, Policy checkpoints and Policy model exports in native format

[sven1977]

Docs update: New Algorithm checkpoints, Policy checkpoints and Policy model exports in native format

Details:
On Algorithm checkpoints:

• All Algorithm checkpoints now use the AIR Checkpoint mechanism. I.e. my_algo.restore([some AIR checkpoint]) works as well as Algorithm.restore([some path to a checkpoint dir]). The checkpoint directory structure will change from:

.
..
checkpoint-[some iter num]

to:

.
..
policies/
    policy_1/
        policy_state.pkl
    policy_2
        policy_state.pkl
checkpoint_version.txt
state.pkl

• Algorithm checkpoints now have a version (e.g. "v0", "v1") stored in the checkpoint dir under "checkpoint_version.txt". This will help keeping checkpoint handling fully backward compatible from Ray 2.0 on. Test cases are introduced in this PR confirming this is and remains the case.
• Algorithm checkpoints now contain a sub-directory ("policies") which has further sub-directories (named after the policies' IDs) that contain the individual policy checkpoints (see below). This allows for easier decomposition and re-assembly of Policies within an Algorithm checkpoint (e.g. restore an Algorithm from a checkpoint, but only with policies A and B, instead of the original A, B, and C, or restoring a Policy instance individually).
• Algorithm gets two new static utilities: from_checkpoint() and from_state(), both of which return new Algorithm objects, given a checkpoint dir or object or a state dict, respectively. I.e.: my_new_algo = Algorithm.from_checkpoint([path to AIR checkpoint OR AIR checkpoint obj]). No original config or other information is needed other than the checkpoint.
• Test cases have been added to keep checkpoint backward compatibility and to test these new utilities and dir structures.

On Policy Checkpoints:

• Policy checkpoints now use the AIR Checkpoint mechanism. I.e. Policy.export_checkpoint() produces an AIR Checkpoint directory with all the policy's state in it.
• Policy gets two new static utilities: from_checkpoint() and from_state(), both of which return new Policy objects, given a Policy checkpoint dir or object or a Policy state, respectively.

On native keras/PyTorch models being part of a Policy checkpoint (optional):

• A new config option: config.checkpointing(checkpoints_contain_native_model_files=True) makes Policies also try to write their NN model as native keras/torch saved model into the given checkpoint directory (under sub-dir "model"). This may still fail (gracefully) in some cases, e.g. for certain TfModelV2 where the keras self.base_model (of the TfModelV2) cannot be discovered easily. This problem will be fully solved by the ongoing RLModule/RLTrainer API efforts.

Why are these changes needed?

Related issue number

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 :(

[maxpumperla]

cool, thanks! @sven1977 as this is a user guide, I think you also want to add this new doc to the panels in user-guides.rst

[maxpumperla]

hmm, maybe we should take a closer look at that gallery, too. For instance, I see that "connectors" are also not in this gallery, although they show up in the TOC (and main navigation of the docs). this should align

[sven1977]

• was already added to user-guides.rst
• @gjoliver on adding connectors docs to list of RLlib user guides (user-guides.rst).

[maxpumperla]

This whole import is too long and consists of 80%+ comments. I think this doesn't read very well on the docs (as a reader I'm not sure if I'm supposed to read all this, it looks as if all those comments are there by "accident"):

[maxpumperla]

basically this import and maybe the last one suffer from this a little, other than that I think this looks great btw.

[sven1977]

I'll split it up and move the comments into the rsv file.

[sven1977]

Left the keras-related comments in the code, but split up the 3 different ways on how to save your models (direct, via policy checkpoint, via algo checkpoint).

[maxpumperla]

Almost there, just the TOC and include file.

[maxpumperla]

not a big deal, but using a :ref: here is more stable and prevents outdated links when moving files etc.

[sven1977]

Done, but can you check, whether I did this correctly. Not sure how sphinx infers the actual html file. E.g. I don't find any serve-rllib-tutorial either (referenced further above this one here).

[maxpumperla]

@sven1977 yeah, I don't like the way Sphinx solves this. essentially, if there's a file called serve-rllib-tutorial.md in the source, relative to the same file, you can use that as doc ref, otherwise (and this is usually much better), you'd just add a .. _serve-rllib-tutorial: tag in rst or (serve-rllib-tutorial)= in markdown in the respective doc and reference that, see here:

https://raw.githubusercontent.com/ray-project/ray/23b3a599b9df8100558c477e94b0b19b1a38ac27/doc/source/serve/tutorials/rllib.md

[maxpumperla]

so in this case, please put .. _rllib-saving-and-loading-algos: at the top of your new doc/source/rllib/rllib-saving-and-loading-algos-and-policies.rst file. then you can reference it as:

:ref:`my reference <rllib-saving-and-loading-algos>`

this should also fix the build

[maxpumperla]

By convention (and I think this is a good one to keep consistent), we have a doc_code folder like here for Tune:

https://github.com/ray-project/ray/tree/master/doc/source/tune/doc_code

which has all the code referenced in docs. You can then basically copy this block in the bazel BUILD file to auto-test all imported snippets. I'm aware that this would likely also happen with rllib/examples but I think it'd be great to keep this code close to doc-sources (which e.g. makes it likely easier for new contributors to understand the structure).

https://github.com/ray-project/ray/blob/master/doc/BUILD#L155-L164

[maxpumperla]

Should be a simple move... apart from that, this is exactly what I had in mind! :D

[sven1977]

Dumb question on this one:
If we currently run all rllib/examples/documentation/* via the CI, how would we reference into the new location (docs/source/rllib/doc_code/*) from the rllib/BUILD file?

[maxpumperla]

we should simply consider those two different CI runs. I didn't realise there was an rllib/examples/documentation folder already. We could either leave it as is, or (in a follow-up PR) move the contents of that folder into doc/source/rllib/doc_code. We can then delete that part from the rllib/BUILD and add the resp. bit to doc/BUILD. that way we don't have to think about the issue you brought up. wdyt?

[maxpumperla]

let's simply skip this discussion for now and fix this later, don't think it's a big deal either way

[maxpumperla]

approving, pending the ref fix