fix delayed reset bug

sail-sg · May 27, 2024 · 1abbeb8 · 1abbeb8
1 parent 9b32dd0
commit 1abbeb8
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/envpool/sokoban/sokoban_envpool.cc b/envpool/sokoban/sokoban_envpool.cc
@@ -142,7 +142,11 @@ void SokobanEnv::Step(const Action& action_dict) {
                         reward_box_ * static_cast<double>(prev_unmatched_boxes -
                                                           unmatched_boxes_) +
                         ((unmatched_boxes_ == 0) ? reward_finished_ : 0.0f);
-  WriteState(static_cast<float>(reward));
+  if (IsDone()) {
+    Reset();
+  } else {
+    WriteState(static_cast<float>(reward));
+  }
 }
 
 constexpr std::array<std::array<uint8_t, 3>, kPlayerOnTarget + 1> kTinyColors{{