Fix vectorized env state persistence and empty episode stats handling

Author: rohanpsingh

rl/algos/ppo.py

@@ -428,7 +428,7 @@ def evaluate(self, env_fn, nets, itr, num_batches=5):
 
         # calculate average evaluation reward
         eval_ep_rewards = [float(i) for batch in eval_batches for i in batch.ep_rewards]
-        avg_eval_ep_rewards = np.mean(eval_ep_rewards)
+        avg_eval_ep_rewards = np.mean(eval_ep_rewards) if eval_ep_rewards else float("nan")
 
         # save checkpoint - saves with suffix and as best if improved
         self.checkpointer.save_if_best(nets, avg_eval_ep_rewards, itr)
@@ -441,6 +441,10 @@ def train(self, env_fn, n_itr):
 
         train_start_time = time.time()
 
+        # Track last known episode stats for iterations where no episodes complete
+        last_mean_ep_reward = float("nan")
+        last_mean_ep_len = float("nan")
+
         obs_mirr, act_mirr = None, None
         if hasattr(env_fn(), "mirror_observation"):
             obs_mirr = env_fn().mirror_clock_observation
@@ -580,9 +584,19 @@ def train(self, env_fn, n_itr):
 
             action_noise = self.policy.stds.data.tolist()
 
+            # Handle empty episode stats (no episodes completed this iteration)
+            if len(batch.ep_rewards) > 0:
+                mean_ep_reward = float(torch.mean(batch.ep_rewards))
+                mean_ep_len = float(torch.mean(batch.ep_lens.float()))
+                last_mean_ep_reward = mean_ep_reward
+                last_mean_ep_len = mean_ep_len
+            else:
+                mean_ep_reward = last_mean_ep_reward
+                mean_ep_len = last_mean_ep_len
+
             sys.stdout.write("-" * 37 + "\n")
-            sys.stdout.write(f"| {'Mean Eprew':>15} | {torch.mean(batch.ep_rewards):>15.5g} |\n")
-            sys.stdout.write(f"| {'Mean Eplen':>15} | {torch.mean(batch.ep_lens.float()):>15.5g} |\n")
+            sys.stdout.write(f"| {'Mean Eprew':>15} | {mean_ep_reward:>15.5g} |\n")
+            sys.stdout.write(f"| {'Mean Eplen':>15} | {mean_ep_len:>15.5g} |\n")
             sys.stdout.write(f"| {'Actor loss':>15} | {np.mean(actor_losses):>15.3g} |\n")
             sys.stdout.write(f"| {'Critic loss':>15} | {np.mean(critic_losses):>15.3g} |\n")
             sys.stdout.write(f"| {'Mirror loss':>15} | {np.mean(mirror_losses):>15.3g} |\n")
@@ -614,8 +628,8 @@ def train(self, env_fn, n_itr):
 
                 eval_ep_lens = [float(i) for b in eval_batches for i in b.ep_lens]
                 eval_ep_rewards = [float(i) for b in eval_batches for i in b.ep_rewards]
-                avg_eval_ep_lens = np.mean(eval_ep_lens)
-                avg_eval_ep_rewards = np.mean(eval_ep_rewards)
+                avg_eval_ep_lens = np.mean(eval_ep_lens) if eval_ep_lens else float("nan")
+                avg_eval_ep_rewards = np.mean(eval_ep_rewards) if eval_ep_rewards else float("nan")
                 print("====EVALUATE EPISODE====")
                 print(
                     f"(Episode length:{avg_eval_ep_lens:.3f}. Reward:{avg_eval_ep_rewards:.3f}. "
@@ -631,8 +645,8 @@ def train(self, env_fn, n_itr):
                 critic_loss=np.mean(critic_losses),
                 mirror_loss=np.mean(mirror_losses),
                 imitation_loss=np.mean(imitation_losses),
-                mean_reward=float(torch.mean(batch.ep_rewards)),
-                mean_ep_len=float(torch.mean(batch.ep_lens.float())),
+                mean_reward=mean_ep_reward,
+                mean_ep_len=mean_ep_len,
                 mean_noise_std=np.mean(action_noise),
                 step=itr,
             )

rl/workers/rollout_worker.py

@@ -80,6 +80,13 @@ def __init__(self, env_fn, policy_template, critic_template, num_envs_per_worker
         self.current_ep_reward = 0.0
         self.current_ep_len = 0
 
+        # Vectorized env state persistence
+        if self.is_vectorized:
+            self.vec_states = None
+            self.vec_traj_lens = None
+            self.vec_ep_reward_accum = None
+            self.vec_ep_len_accum = None
+
     def sync_state(self, policy_state_dict, critic_state_dict, obs_mean, obs_std, iteration_count):
         """Sync all worker state from main process in a single call.
 
@@ -231,22 +238,24 @@ def _sample_vectorized(self, gamma, max_steps, max_traj_len, deterministic):
             PPOBuffer(self.state_dim, self.action_dim, gamma=gamma, size=max_traj_len * 2) for _ in range(num_envs)
         ]
 
-        # Reset all environments
-        states = torch.as_tensor(vec_env.reset_all(), dtype=torch.float)  # [num_envs, obs_dim]
-        traj_lens = torch.zeros(num_envs, dtype=torch.int32)
+        # Initialize or continue from previous state
+        if self.vec_states is None:
+            states = torch.as_tensor(vec_env.reset_all(), dtype=torch.float)
+            traj_lens = torch.zeros(num_envs, dtype=torch.int32)
+            ep_reward_accum = [0.0] * num_envs
+            ep_len_accum = [0] * num_envs
+            if hasattr(policy, "init_hidden_state"):
+                policy.init_hidden_state()
+            if hasattr(critic, "init_hidden_state"):
+                critic.init_hidden_state()
+        else:
+            states = self.vec_states
+            traj_lens = self.vec_traj_lens
+            ep_reward_accum = self.vec_ep_reward_accum
+            ep_len_accum = self.vec_ep_len_accum
+
         ep_lens_per_env = [[] for _ in range(num_envs)]
         ep_rewards_per_env = [[] for _ in range(num_envs)]
-        ep_reward_accum = [0.0] * num_envs
-        ep_len_accum = [0] * num_envs
-
-        # For recurrent policies
-        if hasattr(policy, "init_hidden_state"):
-            # Note: This might need modification for vectorized recurrent policies
-            # Currently resets to same initial state for all envs
-            policy.init_hidden_state()
-
-        if hasattr(critic, "init_hidden_state"):
-            critic.init_hidden_state()
 
         total_steps = 0
         while total_steps < max_steps:
@@ -315,6 +324,12 @@ def _sample_vectorized(self, gamma, max_steps, max_traj_len, deterministic):
                     # Bootstrap with final value (trajectory was truncated, not done)
                     buffers[i].finish_path(last_val=final_value)
 
+        # Persist state for next sample() call
+        self.vec_states = states
+        self.vec_traj_lens = traj_lens
+        self.vec_ep_reward_accum = ep_reward_accum
+        self.vec_ep_len_accum = ep_len_accum
+
         # Flatten episode stats across all envs
         all_ep_lens = [ep_len for per_env in ep_lens_per_env for ep_len in per_env]
         all_ep_rewards = [r for per_env in ep_rewards_per_env for r in per_env]