work

haosulab · Jan 29, 2024 · eb8aece · eb8aece
1 parent a2aa047
commit eb8aece
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/examples/baselines/ppo/ppo.py b/examples/baselines/ppo/ppo.py
@@ -231,7 +231,7 @@ def clip_action(action: torch.Tensor):
             print(f"eval_episodic_return={episodic_return}")
             writer.add_scalar("charts/eval_success_rate", info["success"].float().mean().cpu().numpy(), global_step)
             writer.add_scalar("charts/eval_episodic_return", episodic_return, global_step)
-            writer.add_scalar("charts/eval_episodic_length", info["elapsed_steps"], global_step)
+            writer.add_scalar("charts/eval_episodic_length", info["elapsed_steps"].float().mean().cpu().numpy(), global_step)
         # exit()
         if args.save_model and iteration % args.eval_freq == 1:
             model_path = f"runs/{run_name}/{args.exp_name}_{iteration}.cleanrl_model"
@@ -273,7 +273,7 @@ def clip_action(action: torch.Tensor):
                 print(f"global_step={global_step}, episodic_return={episodic_return}")
                 writer.add_scalar("charts/success_rate", info["success"].float().mean().cpu().numpy(), global_step)
                 writer.add_scalar("charts/episodic_return", episodic_return, global_step)
-                writer.add_scalar("charts/episodic_length", info["elapsed_steps"], global_step)
+                writer.add_scalar("charts/episodic_length", info["elapsed_steps"].float().mean().cpu().numpy(), global_step)
         # bootstrap value if not done
         with torch.no_grad():
             next_value = agent.get_value(next_obs).reshape(1, -1)

diff --git a/mani_skill2/vector/wrappers/gymnasium.py b/mani_skill2/vector/wrappers/gymnasium.py
@@ -55,7 +55,9 @@ def step(
         self.returns += rew
         infos["episode"] = dict(r=self.returns)
         terminations = torch.zeros(self.num_envs, device=self.env.device)
-        if truncations:
+        if (isinstance(truncations, torch.Tensor) and truncations.any()) or (
+            not isinstance(truncations, torch.Tensor) and truncations
+        ):
             infos["episode"]["r"] = self.returns.clone()
             final_obs = obs
             obs, _ = self.reset()