[envpool] Fix explicit episodic-life reset for Atari and Vizdoom (#344)

## Summary
- Problem: explicit `env.reset()` did not fully reset underlying state
for episodic-life envs, so repeated resets could continue the same
game/session instead of starting a fresh episode.
- Scope: preserve the existing `force_reset` signal through the core
async env path, make Atari and Vizdoom honor explicit resets under
`episodic_life`, add regression coverage for both Python-facing APIs,
and keep the branch green against current lint drift.
- Outcome: explicit user resets now route to a real underlying reset for
both env families instead of being treated like episodic-life
continuation.

Fixes #301 by restoring explicit reset semantics for episodic-life envs.

## Technical Details
- Approach: thread the existing `force_reset` bit from `AsyncEnvPool`
into `Env::PreProcess`, then consume it in env-specific reset logic.
- Code pointers:
- `envpool/core/env.h`: stores `force_reset_` on the env instance and
threads it through `EnvStep`/`PreProcess`.
- `envpool/core/async_envpool.h`: forwards `raw_action.force_reset`
instead of collapsing explicit resets into generic reset handling.
- `envpool/atari/atari_env.h`: treats explicit resets as a full
`reset_game()` even when `episodic_life=True`.
- `envpool/vizdoom/vizdoom_env.h`: treats explicit resets as a full
`newEpisode()` even when `episodic_life=True`.
- `envpool/atari/atari_env_test.cc`: adds a C++ regression test for
repeated explicit resets.
- `envpool/atari/atari_envpool_test.py`: adds a gymnasium regression
test matching the user-facing Atari repro.
- `envpool/vizdoom/vizdoom_test.py`: adds a Python regression test for
repeated explicit resets under Vizdoom episodic life.
- `envpool/core/xla_template.h`: fixes current `clang-tidy` / lint
complaints encountered while validating this branch.
- Notes: the Linux `make bazel-test` run on `dev` used `--distdir` with
a locally downloaded `pretrain.tar.gz` and an Atari ROM override to
avoid unrelated external download failures during validation.

## Test Plan
### Automated
- `make bazel-test BAZELOPT="--distdir=/tmp/bazel-distdir-test
--override_repository=atari_roms=/tmp/atari_roms_override_test"` on
`dev`: passed (`30 / 30 tests pass`) before the Vizdoom follow-up.
- issue #301 Atari repro script on `dev` against a patched wheel: passed
(`info["lives"]` stayed at 5 across repeated `reset()`).
- `ruff check envpool/vizdoom/vizdoom_test.py`: passed.
- `ruff format --check envpool/vizdoom/vizdoom_test.py`: passed.
- `clang-format --style=file -n --Werror envpool/vizdoom/vizdoom_env.h`
on `dev`: passed.

### Known Verification Gap
- Targeted `//envpool/vizdoom:vizdoom_test` execution on `dev` currently
crashes in existing test setup with `v_video.cpp:1346 Assertion
'CleanWidth >= 320' failed`, including the pre-existing `test_timelimit`
case. That prevented runtime confirmation of the new Vizdoom regression
test on this box, but the failure reproduces independently of this
patch.

Author: Trinkle23897

envpool/atari/atari_env.h

@@ -151,7 +151,7 @@ class AtariEnv : public Env<AtariEnvSpec> {
   void Reset() override {
     int noop = dist_noop_(gen_) + 1 - static_cast<int>(fire_reset_);
     bool push_all = false;
-    if (!episodic_life_ || env_->game_over() ||
+    if (force_reset_ || !episodic_life_ || env_->game_over() ||
         elapsed_step_ >= max_episode_steps_) {
       env_->reset_game();
       elapsed_step_ = 0;

envpool/atari/atari_env_test.cc

@@ -254,6 +254,26 @@ TEST(AtariEnvTest, EpisodicLife) {
   }
 }
 
+TEST(AtariEnvTest, ExplicitResetWithEpisodicLife) {
+  auto config = atari::AtariEnvSpec::kDefaultConfig;
+  config["num_envs"_] = 1;
+  config["batch_size"_] = 1;
+  config["seed"_] = 42;
+  config["episodic_life"_] = true;
+  config["task"_] = "breakout";
+  atari::AtariEnvSpec spec(config);
+  atari::AtariEnvPool envpool(spec);
+  TArray all_env_ids(Spec<int>({1}));
+  all_env_ids[0] = 0;
+  for (int i = 0; i < 20; ++i) {
+    envpool.Reset(all_env_ids);
+    AtariState state(envpool.Recv());
+    auto lives = state["info:lives"_];
+    int live = lives[0];
+    EXPECT_EQ(live, 5) << "reset #" << i;
+  }
+}
+
 TEST(AtariEnvTest, ZeroDiscountOnLifeLoss) {
   std::srand(std::time(nullptr));
   int batch = 4;

envpool/atari/atari_envpool_test.py

@@ -144,6 +144,22 @@ def test_reset_life(self) -> None:
             self.assertTrue(next_info["lives"][0] > 0)
             self.assertTrue(info["terminated"][0])
 
+    def test_explicit_reset_with_episodic_life_gymnasium(self) -> None:
+        """Issue 301."""
+        env = make_gymnasium(
+            "Breakout-v5",
+            num_envs=1,
+            seed=42,
+            episodic_life=True,
+        )
+        for i in range(20):
+            _, info = env.reset()
+            self.assertEqual(
+                int(info["lives"][0]),
+                5,
+                msg=f"reset #{i} returned lives={info['lives']}",
+            )
+
     def test_partial_step(self) -> None:
         num_envs = 5
         max_episode_steps = 10

envpool/core/async_envpool.h

@@ -124,7 +124,8 @@ class AsyncEnvPool : public EnvPool<typename Env::Spec> {
           int env_id = raw_action.env_id;
           int order = raw_action.order;
           bool reset = raw_action.force_reset || envs_[env_id]->IsDone();
-          envs_[env_id]->EnvStep(state_buffer_queue_.get(), order, reset);
+          envs_[env_id]->EnvStep(state_buffer_queue_.get(), order, reset,
+                                 raw_action.force_reset);
         }
       });
     }

envpool/core/env.h

@@ -65,6 +65,7 @@ class Env {
   EnvSpec spec_;
   int env_id_, seed_;
   std::mt19937 gen_;
+  bool force_reset_{false};
 
  private:
   StateBufferQueue* sbq_;
@@ -159,8 +160,8 @@ class Env {
     }
   }
 
-  void EnvStep(StateBufferQueue* sbq, int order, bool reset) {
-    PreProcess(sbq, order, reset);
+  void EnvStep(StateBufferQueue* sbq, int order, bool reset, bool force_reset) {
+    PreProcess(sbq, order, reset, force_reset);
     if (reset) {
       Reset();
     } else {
@@ -178,9 +179,11 @@ class Env {
   virtual bool IsDone() { throw std::runtime_error("is_done not implemented"); }
 
  protected:
-  void PreProcess(StateBufferQueue* sbq, int order, bool reset) {
+  void PreProcess(StateBufferQueue* sbq, int order, bool reset,
+                  bool force_reset) {
     sbq_ = sbq;
     order_ = order;
+    force_reset_ = force_reset;
     if (reset) {
       current_step_ = 0;
     } else {

envpool/core/xla_template.h

@@ -66,19 +66,19 @@ struct CustomCall {
 
   static xla_ffi::Error ValidateArity(xla_ffi::RemainingArgs args,
                                       xla_ffi::RemainingRets rets) {
-    constexpr std::size_t kExpectedArgs = std::tuple_size_v<In> + 1;
-    constexpr std::size_t kExpectedRets = std::tuple_size_v<Out> + 1;
-    if (args.size() != kExpectedArgs) {
+    constexpr std::size_t expected_args = std::tuple_size_v<In> + 1;
+    constexpr std::size_t expected_rets = std::tuple_size_v<Out> + 1;
+    if (args.size() != expected_args) {
       return xla_ffi::Error::InvalidArgument(
-          "Expected " + std::to_string(kExpectedArgs) + " buffers, got " +
+          "Expected " + std::to_string(expected_args) + " buffers, got " +
           std::to_string(args.size()));
     }
-    if (rets.size() != kExpectedRets) {
+    if (rets.size() != expected_rets) {
       return xla_ffi::Error::InvalidArgument(
-          "Expected " + std::to_string(kExpectedRets) + " results, got " +
+          "Expected " + std::to_string(expected_rets) + " results, got " +
           std::to_string(rets.size()));
     }
-    return xla_ffi::Error();
+    return {};
   }
 
   static xla_ffi::Error PopulateInBuffers(xla_ffi::RemainingArgs args,
@@ -90,7 +90,7 @@ struct CustomCall {
       }
       (*in_arr)[i] = (*buffer).untyped_data();
     }
-    return xla_ffi::Error();
+    return {};
   }
 
   static xla_ffi::Error PopulateOutBuffers(xla_ffi::RemainingRets rets,
@@ -102,7 +102,7 @@ struct CustomCall {
       }
       (*out_arr)[i] = (*buffer)->untyped_data();
     }
-    return xla_ffi::Error();
+    return {};
   }
 
   static xla_ffi::Error CpuExecute(xla_ffi::RemainingArgs args,
@@ -124,7 +124,7 @@ struct CustomCall {
       return err;
     }
     CC::Cpu(*obj, in_arr, out_arr);
-    return xla_ffi::Error();
+    return {};
   }
 
   static xla_ffi::Error GpuExecute(cudaStream_t stream,
@@ -147,7 +147,7 @@ struct CustomCall {
       return err;
     }
     CC::Gpu(*obj, stream, in_arr, out_arr);
-    return xla_ffi::Error();
+    return {};
   }
 
   static auto Specs(Class* obj) {
@@ -177,8 +177,12 @@ struct CustomCall {
                                .RemainingArgs()
                                .RemainingRets()
                                .Attrs<xla_ffi::Dictionary>());
-    return std::make_tuple(py::capsule(reinterpret_cast<void*>(cpu_handler)),
-                           py::capsule(reinterpret_cast<void*>(gpu_handler)));
+    // NOLINTNEXTLINE(bugprone-casting-through-void)
+    auto* cpu_handler_ptr = reinterpret_cast<void*>(cpu_handler);
+    // NOLINTNEXTLINE(bugprone-casting-through-void)
+    auto* gpu_handler_ptr = reinterpret_cast<void*>(gpu_handler);
+    return std::make_tuple(py::capsule(cpu_handler_ptr),
+                           py::capsule(gpu_handler_ptr));
   }
 
   static auto Xla(Class* obj) {

envpool/vizdoom/vizdoom_env.h

@@ -274,7 +274,8 @@ class VizdoomEnv : public Env<VizdoomEnvSpec> {
   bool IsDone() override { return done_; }
 
   void Reset() override {
-    if (dg_->isEpisodeFinished() || elapsed_step_ >= max_episode_steps_) {
+    if (force_reset_ || dg_->isEpisodeFinished() ||
+        elapsed_step_ >= max_episode_steps_) {
       elapsed_step_ = 0;
       if (episode_count_ > 0) {  // NewEpisode at beginning may hang on MAEnv
         if (save_lmp_) {

envpool/vizdoom/vizdoom_test.py

@@ -155,6 +155,64 @@ def test_obs_space(self) -> None:
             == 1 * 4
         )
 
+    def test_explicit_reset_with_episodic_life_gymnasium(self) -> None:
+        env = make_gym(
+            "D1Basic-v1",
+            num_envs=1,
+            seed=42,
+            episodic_life=True,
+            use_combined_action=True,
+        )
+        env.reset()
+        tracked_keys = [
+            "AMMO2",
+            "HEALTH",
+            "HITCOUNT",
+            "KILLCOUNT",
+            "SELECTED_WEAPON_AMMO",
+        ]
+
+        def scalar(info: dict, key: str) -> float:
+            return float(np.asarray(info[key]).reshape(-1)[0])
+
+        action_id = None
+        changed_key = None
+        baseline_value = None
+        changed_value = None
+
+        for candidate in range(env.action_space.n):
+            _, baseline_info = env.reset()
+            for _ in range(64):
+                _, _, terminated, truncated, info = env.step(
+                    np.array([candidate], dtype=int)
+                )
+                for key in tracked_keys:
+                    current = scalar(info, key)
+                    baseline = scalar(baseline_info, key)
+                    if current != baseline:
+                        action_id = candidate
+                        changed_key = key
+                        baseline_value = baseline
+                        changed_value = current
+                        break
+                if changed_key is not None or terminated[0] or truncated[0]:
+                    break
+            if changed_key is not None:
+                break
+
+        assert changed_key is not None
+        assert baseline_value is not None
+        assert changed_value is not None
+        _, reset_info = env.reset()
+        self.assertEqual(
+            scalar(reset_info, changed_key),
+            baseline_value,
+            msg=(
+                f"action={action_id}, key={changed_key}, "
+                f"changed={changed_value}, baseline={baseline_value}"
+            ),
+        )
+
 
 if __name__ == "__main__":
     absltest.main()