Fix progress bar (#30)

# Pull Request

### What change is being made

One line progress bar.

### Why this change is being made

Better user experience.

### Tested

Run evaluator.

GitOrigin-RevId: 7c91a507c1dc1817da990f265a3eb762988e8490

Author: awoll-bdai

exporter/exporter/evaluator.py

@@ -46,8 +46,8 @@ def _print_progress_bar(
     extra_str = " | ".join(extra_info)
 
     print(
-        f"\r{status_emoji} {bar} {step_ctr + 1}/{num_steps} | Failed: {failed_steps} | {extra_str} \n",
-        end="",
+        f"\r{status_emoji} {bar} {step_ctr + 1}/{num_steps} | Failed: {failed_steps} | {extra_str}",
+        end="\n" if not step_export_ok else "",
         flush=True,
     )
 
@@ -61,10 +61,9 @@ def _compare_step_outputs(
     env_outputs: dict[str, torch.Tensor],
     ort_outputs: dict[str, torch.Tensor] | None,
     output_names: list[str],
-    verbose: bool,
     atol: float,
     rtol: float,
-) -> bool:
+) -> tuple[bool, str]:
     """Compare outputs from environment and ONNX model.
 
     Args:
@@ -76,60 +75,64 @@ def _compare_step_outputs(
         env_outputs: Environment outputs.
         ort_outputs: ONNX model outputs (None if session not run yet).
         output_names: Names of output components.
-        verbose: Whether to print comparison results.
         atol: Absolute tolerance.
         rtol: Relative tolerance.
 
     Returns:
-        True if all comparisons passed, False otherwise.
+        A tuple of (is_step_export_ok, message), where `is_step_export_ok` is a boolean indicating whether the step's outputs are close within the specified tolerances, and `message` is a string describing the comparison results, including details of any mismatches and troubleshooting checklists.
     """
     step_export_ok = True
-
-    if verbose:
-        print("\n")
-
-    obs_ok = compare_tensors(
-        vec_a=env_obs.view(1, -1),
-        vec_b=ort_obs.to(env_obs.device).view(1, -1),
+    msg = ""
+
+    # Check observation comparison
+    obs_env = env_obs.view(1, -1)
+    obs_ort = ort_obs.to(env_obs.device).view(1, -1)
+    obs_ok, obs_message = compare_tensors(
+        vec_a=obs_env,
+        vec_b=obs_ort,
         name_a="env",
         name_b="ort",
         vec_name="observation",
         index_names=observation_names,
-        verbose=verbose,
         atol=atol,
         rtol=rtol,
     )
 
-    if not obs_ok and verbose:
-        print("\n📋 Observation Troubleshooting Checklist:")
-        print("  • Verify all observation data sources have corresponding input components")
-        print(
-            "  • Ensure input components reference the same data sources as observation computation"
+    if not obs_ok:
+        msg += "\n\n📋 Observation comparison failed!"
+        msg += f"\n{obs_message}"
+        msg += "\n📋 Observation Troubleshooting Checklist:"
+        msg += "\n  • Verify all observation data sources have corresponding input components"
+        msg += "\n  • Ensure input components reference the same data sources as observation computation"
+        msg += (
+            "\n  • For memory-based observations, confirm memory components are properly registered"
         )
-        print(
-            "  • For memory-based observations, confirm memory components are properly registered"
-        )
-        print("  • Review compute_observation() implementation for data flow correctness")
-        return False
+        msg += "\n  • Review compute_observation() implementation for data flow correctness"
+        return False, msg
 
     step_export_ok = step_export_ok and obs_ok
 
-    actions_ok = compare_tensors(
-        vec_a=env_actions.view(1, -1),
-        vec_b=ort_actions.to(env_actions.device).view(1, -1),
+    # Check actions comparison
+    actions_env = env_actions.view(1, -1)
+    actions_ort = ort_actions.to(env_actions.device).view(1, -1)
+    actions_ok, actions_message = compare_tensors(
+        vec_a=actions_env,
+        vec_b=actions_ort,
         name_a="env",
         name_b="ort",
         vec_name="actions",
-        verbose=verbose,
+        index_names=None,
         atol=atol,
         rtol=rtol,
     )
 
-    if not actions_ok and verbose:
-        print("\n📋 Actions Troubleshooting Checklist:")
-        print("  • Verify actor network matches between env and ONNX")
-        print("  • Ensure action normalizer parameters are correctly exported")
-        return False
+    if not actions_ok:
+        msg += "\n\n🎯 Actions comparison failed!"
+        msg += f"\n{actions_message}"
+        msg += "\n📋 Actions Troubleshooting Checklist:"
+        msg += "\n  • Verify actor network matches between env and ONNX"
+        msg += "\n  • Ensure action normalizer parameters are correctly exported"
+        return False, msg
 
     step_export_ok = step_export_ok and actions_ok
 
@@ -148,28 +151,32 @@ def _compare_step_outputs(
             output_size = env_outputs[name].view(1, -1).shape[1]
             expanded_index_names.extend([f"{name}[{i}]" for i in range(output_size)])
 
-        outputs_ok = compare_tensors(
+        # Check outputs comparison
+        outputs_ok, outputs_message = compare_tensors(
             vec_a=env_outputs_cat,
             vec_b=ort_outputs_cat,
             name_a="env",
             name_b="ort",
             vec_name="outputs",
             index_names=expanded_index_names,
-            verbose=verbose,
             atol=atol,
             rtol=rtol,
         )
 
-        if not outputs_ok and verbose:
-            print("\n📋 Outputs Troubleshooting Checklist:")
-            print("  • Verify output components are registered for all expected outputs")
-            print("  • Ensure the correct processed actions are included in memory")
-            print("  • Review process_action() and apply_action() implementations for consistency")
-            return False
+        if not outputs_ok:
+            msg += "\n\n📊 Outputs comparison failed!"
+            msg += f"\n{outputs_message}"
+            msg += "\n📋 Outputs Troubleshooting Checklist:"
+            msg += "\n  • Verify output components are registered for all expected outputs"
+            msg += "\n  • Ensure the correct processed actions are included in memory"
+            msg += (
+                "\n  • Review process_action() and apply_action() implementations for consistency"
+            )
+            return False, msg
 
         step_export_ok = step_export_ok and outputs_ok
 
-    return step_export_ok
+    return step_export_ok, msg
 
 
 def evaluate(
@@ -215,6 +222,10 @@ def evaluate(
 
     obs = observations.clone() if observations is not None else env.observations_reset()
 
+    # Print ONNX graph structure if verbose
+    if verbose:
+        session_wrapper.print_graph()
+
     step_ctr = 0
     export_ok = True
     failed_steps = 0
@@ -313,7 +324,7 @@ def reset():
         }
 
         # Compare outputs from environment and ONNX model.
-        step_export_ok = _compare_step_outputs(
+        step_export_ok, msg = _compare_step_outputs(
             env_obs=obs,
             ort_obs=ort_observations,
             env_actions=env_actions,
@@ -322,7 +333,6 @@ def reset():
             ort_outputs=ort_outputs,
             observation_names=env.get_observation_names(),
             output_names=context_manager.get_output_names(),
-            verbose=verbose,
             atol=atol,
             rtol=rtol,
         )
@@ -333,6 +343,10 @@ def reset():
 
         # Display progress bar.
         if verbose:
+            if step_ctr == 0:
+                print("\n\nStarting evaluation...")
+            if not step_export_ok:
+                print(msg)
             _print_progress_bar(
                 step_ctr=step_ctr,
                 num_steps=num_steps,
@@ -347,5 +361,6 @@ def reset():
             input()
 
         step_ctr += 1
-
+    if verbose:
+        print("\nEvaluation complete.")
     return export_ok, next_obs

exporter/exporter/session_wrapper.py

@@ -3,8 +3,10 @@
 import pathlib
 
 import numpy as np
+import onnx
 import onnxruntime as ort
 import torch
+from onnx import helper
 
 from exporter.utils.paths import prepare_onnx_paths
 
@@ -108,3 +110,14 @@ def get_output_value(self, output_name: str):
     def reset(self):
         """Reset the internal results to zeros to avoid stale data at environment reset."""
         self._results = [np.zeros_like(output) for output in self._results]
+
+    def print_graph(self) -> None:
+        """Print the ONNX graph structure in a readable format."""
+        try:
+            print("\n📊 ONNX Graph Structure:")
+            print("=" * 50)
+            onnx_model = onnx.load(str(self._onnx_file_path))
+            print(helper.printable_graph(onnx_model.graph))
+            print("=" * 50)
+        except Exception as e:
+            print(f"⚠️  Could not load ONNX graph: {e}")

exporter/exporter/utils/math.py

@@ -13,8 +13,7 @@ def compare_tensors(
     name_b: str = None,
     atol: float = 1.0e-5,
     rtol: float = 1.0e-5,
-    verbose: bool = False,
-) -> bool:
+) -> tuple[bool, str]:
     """Compare two tensors, and print a message showing the differences.
 
     Args:
@@ -23,18 +22,19 @@ def compare_tensors(
         vec_name: A name identifying the compared tensors.
         index_names: A list of names, one for each element of the input tensors.
         name_a: A name identifying the source of the first input tensor.
-        name_a: A name identifying the source of the second input tensor.
+        name_b: A name identifying the source of the second input tensor.
         atol : Absolute tolerance.
         rtol : Relative tolerance.
-        verbose: If true, print information about the tensor difference to console.
+
+    Returns:
+        A tuple of (is_close, message), where `is_close` is a boolean indicating whether the tensors are close within the specified tolerances, and `message` is a string describing the comparison results, including details of any mismatches.
     """
     is_close = torch.isclose(vec_a, vec_b, atol=atol, rtol=rtol)
 
     msg = f"Comparing {vec_name}: "
     if torch.all(is_close):
         msg += "\033[92mall elements are close.\033[0m"
-        print(msg)
-        return True
+        return True, msg
 
     mismatched_indices = (torch.logical_not(is_close)).nonzero(as_tuple=False)
 
@@ -52,7 +52,4 @@ def compare_tensors(
             f"{name_b}={val_b:+.5f}\t"
             f"Δ={np.abs(val_a - val_b):.5f}"
         )
-
-    if verbose:
-        print(msg)
-    return False
+    return False, msg