scan over different RL envs for a multitask experiment

[yooceii]

Hi

I am working on a multitask RL experiment.
I would like to change the follow for loop into a lax scan given the later iteration's runner_state depends on the priori one.
And the current loop is fairly slow with a large iterations.

envs, envs_params = list(
      zip(
          *jax.tree.map(
              lambda x: make_env(x),
              env_names,
          )
      )
  )
for _ in range(iterations):
    for env in envs:
        runner_state = update_step(runner_state, env) # similar to https://github.com/luchris429/purejaxrl/blob/9d803f91737ce0925dc0f3054b9e98a580033162/purejaxrl/ppo.py#L111-L270

I have tried to put the envs in the runner_state and scan over env idx, but raises This concrete value was not available in Python because it depends on the value of the argument idx.

Wonder if any one has idea on how to achieve that.

Thanks.
Chang