Update env config in workers

Hi there,

I’m trying to run some experiments a trained model. This includes running an evaluation then updating an env config variable and running the evaluation again. At the moment my code is pretty hacky and incurs a lot of overheads as I couldn’t figure out how to update the env’s config without recreating the trainer.

 for score, hands in handscores.items():        
        config["env_config"].update({"starting_hands": {0: hands}})

        trainer = A3CTrainer(env="yaniv", config=config)
        trainer.restore(args.checkpoint)

        metrics = trainer._evaluate()
        metrics["evaluation"].pop("hist_stats")

        print(pretty_print(metrics))

        trainer.cleanup()

How can I update the config of all the workers?

Cheers,

Rory

Hey @Rory ,
you could try forcing the workers to re-create their envs. Currently there is no API for updating an existing env object on-the-fly.

Like so:

def update_env_fn(worker, worker_idx):
    # Create the updated EnvContext object with the correct worker idx and number of workers.
    new_env_context = EnvContext(
        [your new config], worker_idx, num_workers=[your num workers])
    # Override the Worker's env with a new one.
    worker.env = worker.env_creator(new_env_context)

trainer.workers.foreach_worker_with_index(update_env_fn)

Also, somehow you would have to get your new config dict into the update_env_fn. Maybe w/ a global var or via functools.partial?

Hi @sven1977 ,

Thanks for pointing me in the right direction. Took me a while to figure out but you the code you posted wouldn’t work for the rollout workers, since the sampler is created in the constructor and has a reference to the env. So by updating the env property on the worker it doesn’t update the env in the sampler. Found that each worker has a foreach_env member function so I did this instead:

def make_update_env_fn(env_conf):
    def update_env_conf(env):
        env.config.update(env_conf)
        env.game.configure(env.config)
        
    def update_env_fn(worker):
        worker.foreach_env(update_env_conf)

    return update_env_fn

trainer.evaluation_workers.foreach_worker(
            make_update_env_fn(config["env_config"])
        )

Thanks for the help! Muuuuch faster now :slight_smile:

4 Likes

Awesome! Ah yes, makes sense! Thanks for sharing this here @Rory !