RL Benchmark Demo Planning

[nickjalbert]

Below we track our working demo of the AgentOS Benchmark, built on PCS for RL researchers.

This is a 30min meeting with 10min live code demo.

Goals

• In very few lines of simple intuitive code...
• Download, run, and train a sophisticated agent (Ilya's, Acme, etc.)
• Demonstrate automatic tracking (including MLflow UI)
• Share/publish results
• Show off benefits of public benchmark registry

General outline for the whole demo talk

• introduce the problem we are trying to solve
• show how to install use AOS/PCS
• Create an AgentRun from an agent that is an existing github repo: pytorch-a2c-ppo-acktr-gail
  • Show how we both train, evaluate, and track lineage of that agent's runs and their artifacts (trained DL models)
  • Demo follow hyperlinks in the WebRegistry app view of a Run to its dependencies, including other Run objects that are input dependencies.
• Show off built in experiment tracking features (including MLflow UI & CLI access to tracking)
• Publish the run results to benchmark website, including its recursive dependencies (other runs, RunCommand, Components, Repos, etc.)
• View Registry website and see the run in the context of other benchmark runs (e.g. new Ilya PPO run vs an SB3 PPO run already in the registry)
• Show downloading and re-running a given run locally
  • stretch: have the user try to install AgentOS and reproduce (i.e. re-run) the run we just published

The current script for the code demo

The code below demonstrates the following functionality:

1. infer a registry from SB3
2. create a agent component from that registry
3. run that agent using agentos
4. use the run_command generated from the initial run to do a re-run.
5. publish the run to different types of registries (local file and WebRegistry)
6. read in the run from those different registries, and re-run it

from agentos.repo import Repo
from agentos.registry import Registry
from agentos.component import Component
from agentos.argument_set import ArgumentSet
from agentos.registry import WebRegistry
from agentos.component_run import ComponentRun


sb3_repo = Repo.from_github("DLR-RM", "stable-baselines3")
sb3_reg = Registry.from_repo(sb3_repo)

# If you're doing dev on the demo, you can avoid re-running the 2min infer command above
# by caching the inferred registry to local disk once, and then reading it in and using it
# for subsequent testing runs.
#sb3_reg.to_yaml("/tmp/sb3_registry_inferred.yaml")
#sb3_reg = Registry.from_yaml("/tmp/sb3_registry_inferred.yaml")

sb3_comp = Component.from_registry(sb3_reg, "module:stable_baselines3____init__.py")
sb3_comp.class_name = "PPO"
sb3_comp.instantiate = True
arg_set = ArgumentSet({
        "module:stable_baselines3____init__.py": {
            "__init__": {
                "policy": 'MlpPolicy', "env": "CartPole-v1"
            },
            "learn": {
                "total_timesteps": 10000
            }
        }
    })
run = sb3_comp.run_with_arg_set("learn", arg_set, log_return_value=False)

# Write a run to your local WebRegistry, then read it back in.
wr = WebRegistry("http://localhost:8000/api/v1")
run.to_registry(wr)
run_from_web_reg = ComponentRun.from_registry(wr, run.identifier)

# Try creating a new run from the run_command loaded in from the WebRegistry
new_run = run.run_command.run()

# Write the new run to a file, read it back in, and run it.
new_run.to_registry().to_yaml("/tmp/run_sb3.yaml")
loaded_run_reg = Registry.from_yaml("/tmp/run_sb3.yaml")
loaded_run = ComponentRun.from_registry(loaded_run_reg, new_run.identifier)
loaded_run.run_command.run()

# Make sure we can write a run that was read in from a file registry
# out to a WebRegistry.
#loaded_run.to_registry(wr)
#ComponentRun.from_registry(wr, loaded_run.identifier)

A more ambitious demo that would include more Agent level semantics wrapping lower PCS APIs

>>> from agentos import AgentRepo
>>> repo = AgentRepo.from_github('ikostrikov', 'pytorch-a2c-ppo-acktr-gail')
>>> repo.list_agents()  # looks for components.yaml, also has optional `registry_file: str` parameter. If no registry file exists, you can infer a registry using `Registry.from_repo`
['ppo', 'a2c_acktr']
>>> ppo_agent = repo.get_agent('ppo')
>>> ppo_agent.list_hyperparameters()
['lr', 'clip-param', 'value-loss-coef', 'num-processes', 'num-steps', 'num-mini-batch', 'log-interval', 'use-linear-lr-decay', 'entropy-coef']
>>> ppo_agent.list_versions()
['new', 'new_pong_optimized', 'trained_pong_optimized']
>>> ppo_agent.enable_version('trained_pong_optimized')
>>> env = repo.get_environment('PongNoFrameskip-v4')
>>> initial_obs = env.reset()
>>> action = ppo_agent.act(initial_obs)
>>> obs, reward, done, info = env.step(action)

Notes:

• This is a programmatic version of the command specified here
• Use components.yaml file to specify the agents within a repo and the different versions of a particular agent. Versions can include parameters and artifacts
• Simple interface to get a specific agent and a specific version of that agent
• This agent can then be used to directly on the REPL or in the notebook
• Updates to the agent can be published to AOS repo
• This can be reduced a one-liner in a repo's quickstart to make it extremely easy for developers to test/update/modify:

  agent = Repo.from_github('ikostrikov', 'pytorch-a2c-ppo-acktr-gail').get_agent('ppo').enable_version('trained_pong_optimized')
  

[andyk]

I like this focus on the RL developer persona a lot.