[Feature Request] A simple (and effective?) way to support cherry-picked env reset in `xla` mode

[bkkgbkjb]

Motivation

AFAIK (and also in #194), currently it's unable to cherry-pick terminated envs for reset in xla mode as:

1. envs.reset() method doesn't accept env_ids arguments, unlike sync mode. And even if it does:
2. (step_env() or send() seem to accept env_ids arguments,) but it's hard to generate dynamic-shaped env_ids array in xla mode and jax. (I've made some preliminary exps. to verify this)

The problem resulted from this, would be incorrect transitions to be appeared: (term_state, any_action, rew -> init_state)
as also pointed out in #194

Solution

I think we could try adding masked_env_ids in reset() and send() methods

SO we could do something like this

obs, _ = envs.reset()
handle, recv, send, step = envs.xla()

while True:

    handle, (obs, rew, terms, truncs, info) = step(handle, some_acts)

    # proposed masked auto-resetting
    auto_reset_masks = jnp.logical_or(terms, truncs)
    _obs = env.reset(masked_env_ids = auto_reset_masks)

    obs = jnp.where(auto_reset_masks, _obs, obs) 

in xla mode

masked_env_ids has static shape of env_nums and would only reset envs of True in masks and return dummy obs for False-ed envs.

Alternative Methods

Currently, I'm working-around this inconvenience by overwriting the wrong transitions by previous correct ones. This shall not make a significant difference to general algorithms.

But if the proposed solution is correct, I think it's better to have it for elegance.

Additional context

Unfortunately, I'm not an expert in C++ and I'm not sure if the proposed solution, despite simple, would work as expected.
But based on my understanding, this shall be implementable so long as we perform it in C++ processes.

Checklist

• [BUG] Vectorized Environment Autoreset Incompatible with openai/baselines' API #194 is somehow related to this, but I think the original issue is more about the correct behaviors of auto-reset, while this focuses a simple (but effective) way to resolve a minor issue in current impl.

[JesseFarebro]

This would be great to have, as performing exact evaluation like this: #113 (comment) isn't possible to jit.

[Trinkle23897]

Thanks for the request. One clarification: the current EnvPool auto-reset does not reset all envs. It is decided per env when that env receives the next action: force_reset || envs_[env_id]->IsDone(), so only the env that is already done is reset.

The behavior described here is mostly about timing/API shape in XLA mode. With the current XLA send/step API, after a terminal state is returned, reset happens on the next send/step for that same env, and the action supplied from the terminal observation is discarded. So a terminal_state -> initial_state record is an expected consequence of EnvPool's auto-reset semantics; rollout code should mask done transitions accordingly.

A mask-shaped reset custom call could make some JAX/XLA exact-evaluation loops cleaner because it avoids dynamic-shaped env_ids, but it would be a new API/semantic surface rather than a fix for current auto-reset. We are not planning to add it at the moment, so I'll close this as not planned. Happy to reconsider or review a PR if someone wants to prototype it.