Print per-epoch episode statistics to stdout in addition to TensorBoard logging.
| 98 | |
| 99 | |
| 100 | class DebugIsaacAlgoObserver(IsaacAlgoObserver): |
| 101 | """Print per-epoch episode statistics to stdout in addition to TensorBoard logging.""" |
| 102 | |
| 103 | def _summarize_episode_infos(self) -> dict[str, float]: |
| 104 | if not self.ep_infos: |
| 105 | return {} |
| 106 | |
| 107 | summary: dict[str, float] = {} |
| 108 | for key in self.ep_infos[0]: |
| 109 | info_tensor = torch.tensor([], device=self.algo.device) |
| 110 | for ep_info in self.ep_infos: |
| 111 | value = ep_info[key] |
| 112 | if not isinstance(value, torch.Tensor): |
| 113 | value = torch.tensor([value], dtype=torch.float32, device=self.algo.device) |
| 114 | else: |
| 115 | value = value.to(self.algo.device) |
| 116 | if value.ndim == 0: |
| 117 | value = value.unsqueeze(0) |
| 118 | info_tensor = torch.cat((info_tensor, value)) |
| 119 | summary[key] = torch.mean(info_tensor).item() |
| 120 | return summary |
| 121 | |
| 122 | def after_print_stats(self, frame, epoch_num, total_time): |
| 123 | ep_summary = self._summarize_episode_infos() |
| 124 | direct_summary: dict[str, float] = {} |
| 125 | for key, value in self.direct_info.items(): |
| 126 | if isinstance(value, torch.Tensor): |
| 127 | direct_summary[key] = value.item() |
| 128 | else: |
| 129 | direct_summary[key] = float(value) |
| 130 | |
| 131 | super().after_print_stats(frame, epoch_num, total_time) |
| 132 | |
| 133 | print("\n" + "*" * 100) |
| 134 | print(f"[EPOCH {epoch_num}] frame={frame} total_time={total_time:.2f}s") |
| 135 | if ep_summary: |
| 136 | print("[EPISODE]") |
| 137 | for key in sorted(ep_summary): |
| 138 | print(f" {key}: {ep_summary[key]:.6f}") |
| 139 | else: |
| 140 | print("[EPISODE] no episodic stats collected this epoch") |
| 141 | if direct_summary: |
| 142 | print("[DIRECT]") |
| 143 | for key in sorted(direct_summary): |
| 144 | print(f" {key}: {direct_summary[key]:.6f}") |
| 145 | print("*" * 100) |
| 146 | |
| 147 | |
| 148 | def _print_env_debug(env, env_cfg, agent_cfg, log_root_path: str, log_dir: str) -> None: |