Pin framework to tf in AIR rl offline trainer example

[avnishn]

something broken about the way that torch ppo writes out to json

when using the output writer, torch ppo writes out sample batches in the incorrect format:

“value”: {”type”: “MultiAgentBatch”...

instead of just

{”type”: “MultiAgentBatch”...

Useful debug script:
https://gist.github.com/avnishn/c8599a9947edcd9934421713ed82f175

def generate_offline_data(path: str):
    print(f"Generating offline data for training at {path}")
    trainer = RLTrainer(
        algorithm="PPO",
        run_config=RunConfig(stop={"timesteps_total": 5000}),
        config={
            "env": "CartPole-v1",
            "output": "dataset",
            "output_config": {
                "format": "json",
                "path": path,
                "max_num_samples_per_file": 1,
            },
            "batch_mode": "complete_episodes",
            "framework": "tf"
        },
    )
    trainer.fit()

def train_rl_bc_offline(path: str, num_workers: int, use_gpu: bool = False) -> Result:
    print("Starting offline training")
    dataset = ray.data.read_json(
        path, parallelism=num_workers, ray_remote_args={"num_cpus": 1}
    )

    trainer = RLTrainer(
        run_config=RunConfig(stop={"training_iteration": 5}),
        scaling_config=ScalingConfig(num_workers=num_workers, use_gpu=use_gpu),
        datasets={"train": dataset},
        algorithm=BC,
        config={
            "env": "CartPole-v1",
            "framework": "tf",
            "evaluation_num_workers": 1,
            "evaluation_interval": 1,
            "evaluation_config": {"input": "sampler"},
            "framework": "tf"
        },
    )

    # Todo (krfricke/xwjiang): Enable checkpoint config in RunConfig
    # result = trainer.fit()
    tuner = Tuner(
        trainer,
        _tuner_kwargs={"checkpoint_at_end": True},
    )
    result = tuner.fit()[0]
    return result

ray.init(num_cpus=8)

path = "/tmp/out"
generate_offline_data(path)

result = train_rl_bc_offline(path=path, num_workers=2, use_gpu=False)

Should work with torch but doesn’t because of the way that sample batches are created.

Signed-off-by: Avnish avnishnarayan@gmail.com

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 added any new APIs to the API Reference. For example, if I added a
    method in Tune, I've added it in doc/source/tune/api/ under the
    corresponding .rst file.
• 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 :(

[avnishn]

corresponding issue:
#33751

[amogkam]

can we fix the underlying issue instead?

[avnishn]

coming in a separate pr @amogkam. I can't fix it before the release cut. This will unblock for now

[amogkam]

@gjoliver @matthewdeng please review

[matthewdeng]

We should keep the output?

[avnishn]

sure thing

[avnishn]

ok all done.

[gjoliver]

wait the branch cut has already passed, do we still need this?

[avnishn]

branch cut is tonight at midnight. I was asked to make this fix since this test was red on master but I'll defer to you guys instead to determine whether or not this needs to be merged.