feat: add command-line arguments for sample episodes and save plot directory; update MSE histogram plotting

Signed-off-by: Dongyun Kim <[email protected]>

Author: DongyunRobotis

physical_ai_server/physical_ai_server/evaluation/evaluation_manager.py

@@ -274,6 +274,19 @@ def main():
         required=True,
         help='The path to the policy file to be evaluated.'
     )
+    parser.add_argument(
+        '--sample_episodes',
+        type=int,
+        nargs='+',
+        default=None,
+        help='A list of episode indices to evaluate. If not provided, all episodes are evaluated.'
+    )
+    parser.add_argument(
+        '--save_plot_dir',
+        type=str,
+        default='./plots',
+        help='The directory to save evaluation plots.'
+    )
     args = parser.parse_args()
 
     # Initialize managers
@@ -297,10 +310,10 @@ def main():
             results = evaluation_manager.evaluate_policy_on_dataset(
                 inference_manager=inference_manager,
                 dataset=dataset,
-                sample_episodes=[1, 3, 5],
+                sample_episodes=args.sample_episodes,
                 plot_episodes=True,
                 plot_summary=True,
-                save_plot_dir='./plots'
+                save_plot_dir=args.save_plot_dir
             )
             print(f'Evaluation results: {results}')
 

physical_ai_server/physical_ai_server/evaluation/visualization_manager.py

@@ -25,7 +25,7 @@
 
 class VisualizationManager:
 
-    def __init__(self, default_save_dir: str = '/root/.cache'):
+    def __init__(self, default_save_dir: str = os.path.expanduser('~/.cache')):
         self.default_save_dir = default_save_dir
         os.makedirs(self.default_save_dir, exist_ok=True)
 
@@ -212,6 +212,15 @@ def plot_episode_mse_distribution(
         # Create figure and axis
         _, ax = plt.subplots(figsize=(10, 6))
 
+        # Plot histogram of MSE values
+        ax.hist(
+            episode_mses,
+            bins=20,
+            color='blue',
+            alpha=0.7,
+            edgecolor='black',
+            label='MSE Distribution')
+
         # Add vertical lines for mean and median
         mean_mse = np.mean(episode_mses)
         median_mse = np.median(episode_mses)