Error between 'pbt-cso' agent and 'envpool' env

[RichieZou]

Hello,

Thank you for sharing so versatile tool.

I encounted a error when I run code in CLI:
(evorl) richie@ubuntu2204:~/evorl$ python -m evorl.train agent=meta/pbt-cso-paramppo env=envpool/gym/walker2d_v3

and the error shows up:

Error executing job with overrides: ['agent=meta/pbt-cso-paramppo', 'env=envpool/gym/walker2d_v3']
Traceback (most recent call last):
  File "/home/richie/anaconda3/envs/evorl/lib/python3.10/runpy.py", line 196, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/home/richie/anaconda3/envs/evorl/lib/python3.10/runpy.py", line 86, in _run_code
    exec(code, run_globals)
  File "/home/richie/evorl/evorl/train.py", line 82, in <module>
    train()
  File "/home/richie/anaconda3/envs/evorl/lib/python3.10/site-packages/hydra/main.py", line 94, in decorated_main
    _run_hydra(
  File "/home/richie/anaconda3/envs/evorl/lib/python3.10/site-packages/hydra/_internal/utils.py", line 394, in _run_hydra
    _run_app(
  File "/home/richie/anaconda3/envs/evorl/lib/python3.10/site-packages/hydra/_internal/utils.py", line 457, in _run_app
    run_and_report(
  File "/home/richie/anaconda3/envs/evorl/lib/python3.10/site-packages/hydra/_internal/utils.py", line 220, in run_and_report
    return func()
  File "/home/richie/anaconda3/envs/evorl/lib/python3.10/site-packages/hydra/_internal/utils.py", line 458, in <lambda>
    lambda: hydra.run(
  File "/home/richie/anaconda3/envs/evorl/lib/python3.10/site-packages/hydra/_internal/hydra.py", line 119, in run
    ret = run_job(
  File "/home/richie/anaconda3/envs/evorl/lib/python3.10/site-packages/hydra/core/utils.py", line 186, in run_job
    ret.return_value = task_function(task_cfg)
  File "/home/richie/evorl/evorl/train.py", line 73, in train
    state = workflow.learn(state)
  File "/home/richie/evorl/evorl/algorithms/meta/pbt_base.py", line 407, in learn
    train_metrics, state = self.step(state)
  File "/home/richie/evorl/evorl/algorithms/meta/pbt_base.py", line 247, in step
    pop_eval_metrics, pop_workflow_state = eval_fn(pop_workflow_state)
  File "/home/richie/evorl/evorl/distributed/sharding.py", line 20, in shmap_f
    return jax.vmap(fn)(*args)
  File "/home/richie/evorl/evorl/workflows/rl_workflow.py", line 152, in evaluate
    raw_eval_metrics = self.evaluator.evaluate(
  File "/home/richie/evorl/evorl/evaluators/evaluator.py", line 80, in evaluate
    _, (episode_returns, episode_lengths) = jax.lax.scan(
  File "/home/richie/evorl/evorl/evaluators/evaluator.py", line 50, in _evaluate_fn
    episode_metrics, env_state = fast_eval_rollout_episode(
  File "/home/richie/evorl/evorl/rollout.py", line 364, in fast_eval_rollout_episode
    env_state, _, metrics = jax.lax.while_loop(
jax._src.source_info_util.JaxStackTraceBeforeTransformation: Exception: Unordered IO effects not supported in while_loop with batched predicate

The preceding stack trace is the source of the JAX operation that, once transformed by JAX, triggered the following exception.

--------------------

The above exception was the direct cause of the following exception:

jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/richie/evorl/evorl/train.py", line 76, in train
    raise e
  File "/home/richie/evorl/evorl/train.py", line 73, in train
    state = workflow.learn(state)
  File "/home/richie/evorl/evorl/algorithms/meta/pbt_base.py", line 407, in learn
    train_metrics, state = self.step(state)
  File "/home/richie/evorl/evorl/algorithms/meta/pbt_base.py", line 247, in step
    pop_eval_metrics, pop_workflow_state = eval_fn(pop_workflow_state)
  File "/home/richie/evorl/evorl/distributed/sharding.py", line 20, in shmap_f
    return jax.vmap(fn)(*args)
Exception: Unordered IO effects not supported in while_loop with batched predicate

Set the environment variable HYDRA_FULL_ERROR=1 for a complete stack trace.

Maybe this is the bug between pbt-cso agent and envpool env, the other two envs ( brax and gymnax ) is accompatible with pbt-cso agent.

Many thanks.

[ZaberKo]

Currently, the support of cpu-based envs by envpool is limited. By now, we support normal RL and EC pipelines for these envs. Since we offer full jit compatible pipelines, handling these CPU-based environments can be tricky through multiple vamp(). We are trying to explore some internal mechanisms inside JAX to handle it.

[artiko88]

Nice! By normal RL and EC pipelines - do you mean evolutionary algorithms?