cleaned learning example, added animation

Author: simeon-ned

examples/pendulum/02_simulation_error.py renamed to _draft/02_simulation_error.py

@@ -12,6 +12,7 @@
 import os
 import optax
 from mujoco.mjx._src.types import IntegratorType
+from _plotting_utils import plot_simulation_errors
 
 # SHOULD WE MOVE THIS IN TO MODULE INIT?
 xla_flags = os.environ.get("XLA_FLAGS", "")
@@ -40,7 +41,7 @@ def parameters_map(parameters: jnp.ndarray, model: mjx.Model) -> mjx.Model:
 key = jax.random.PRNGKey(0)
 
 # Load the model
-MJCF_PATH = "models/pendulum.xml"
+MJCF_PATH = "models/pendulum_estimated.xml"
 model = mujoco.MjModel.from_xml_path(MJCF_PATH)
 data = mujoco.MjData(model)
 model.opt.integrator = IntegratorType.EULER
@@ -112,45 +113,19 @@ def parameters_map(parameters: jnp.ndarray, model: mjx.Model) -> mjx.Model:
 
 predicted_terminal_points = np.array(batched_states_trajectories)[:, -1, :]
 batched_states_trajectories = np.array(batched_states_trajectories).reshape(N_INTERVALS * HORIZON, 2)
-# Plotting simulation results for batсhed state trajectories
-plt.figure(figsize=(10, 5))
-
-plt.subplot(2, 2, 1)
-plt.plot(timespan, angle, label="Actual Angle", color="black", linestyle="dashed", linewidth=2)
-plt.plot(timespan, batched_states_trajectories[:, 0], alpha=0.5, color="blue", label="Simulated Angle")
-plt.plot(timespan, angle, label="Actual Angle", color="black", linestyle="dashed", linewidth=2)
-plt.plot(timespan[HORIZON + 1 :][::HORIZON], predicted_terminal_points[:-1, 0], "ob")
-plt.plot(timespan[HORIZON + 1 :][::HORIZON], interval_terminal_states[:, 0], "or")
-plt.ylabel("Angle (rad)")
-plt.grid(color="black", linestyle="--", linewidth=1.0, alpha=0.4)
-plt.legend()
-plt.title("Pendulum Dynamics - Bathed State Trajectories")
-
-plt.subplot(2, 2, 3)
-plt.plot(timespan, velocity, label="Actual Velocity", color="black", linestyle="dashed", linewidth=2)
-plt.plot(timespan[HORIZON + 1 :][::HORIZON], predicted_terminal_points[:-1, 1], "ob")
-plt.plot(timespan[HORIZON + 1 :][::HORIZON], interval_terminal_states[:, 1], "or")
-plt.plot(timespan, batched_states_trajectories[:, 1], alpha=0.5, color="blue", label="Simulated Velocity")
-plt.xlabel("Time (s)")
-plt.ylabel("Velocity (rad/s)")
-plt.grid(color="black", linestyle="--", linewidth=1.0, alpha=0.4)
-plt.legend()
-
-# Add phase portrait
-plt.subplot(1, 2, 2)
-plt.plot(angle, velocity, label="Actual", color="black", linestyle="dashed", linewidth=2)
-plt.plot(
-    batched_states_trajectories[:, 0], batched_states_trajectories[:, 1], alpha=0.5, color="blue", label="Simulated"
+
+# Replace the plotting section with:
+plot_simulation_errors(
+    timespan, 
+    angle, 
+    velocity, 
+    batched_states_trajectories, 
+    predicted_terminal_points, 
+    interval_terminal_states, 
+    HORIZON,
+    save_path="plots/simulation_error.png",
+    show=True
 )
-plt.plot(predicted_terminal_points[:-1, 0], predicted_terminal_points[:-1, 1], "ob")
-plt.plot(interval_terminal_states[:, 0], interval_terminal_states[:, 1], "or")
-plt.xlabel("Angle (rad)")
-plt.ylabel("Angular Velocity (rad/s)")
-plt.title("Phase Portrait")
-plt.grid(color="black", linestyle="--", linewidth=1.0, alpha=0.4)
-plt.legend()
-
-plt.tight_layout()
-plt.show()
+
 # TODO:
 # Optimization

examples/pendulum/TODO

@@ -1 +1,7 @@
-☐ Jupyter notebook with pendulum learning
+☐ Jupyter notebook with pendulum learning
+    ☐ Parametric model
+    ☐ Simulation error 
+    ☐ Residuals 
+    ☐ Learning
+    ☐ Media animation 
+    

examples/pendulum/_plotting_utils.py

@@ -0,0 +1,206 @@
+import matplotlib.pyplot as plt
+import numpy as np
+import os
+from typing import List, Tuple, Union
+
+def setup_plot(nrows: int = 2, ncols: int = 2, figsize: Tuple[int, int] = (10, 10)) -> Tuple[plt.Figure, Union[plt.Axes, List[plt.Axes]]]:
+    """
+    Set up a matplotlib figure with the specified number of rows and columns.
+
+    Args:
+        nrows (int): Number of rows in the subplot grid.
+        ncols (int): Number of columns in the subplot grid.
+        figsize (Tuple[int, int]): Figure size in inches (width, height).
+
+    Returns:
+        Tuple[plt.Figure, Union[plt.Axes, List[plt.Axes]]]: Figure and axes objects.
+    """
+    fig, axes = plt.subplots(nrows, ncols, figsize=figsize)
+    if nrows * ncols == 1:
+        axes = [axes]
+    elif nrows == 1 or ncols == 1:
+        axes = axes.flatten()
+    return fig, axes
+
+def plot_state(ax: plt.Axes, timespan: np.ndarray, actual: np.ndarray, simulated: np.ndarray, label: str, color: str = 'blue', alpha: float = 0.5) -> None:
+    """
+    Plot actual and simulated state data on the given axes.
+
+    Args:
+        ax (plt.Axes): The matplotlib axes to plot on.
+        timespan (np.ndarray): Array of time points.
+        actual (np.ndarray): Array of actual state values.
+        simulated (np.ndarray): Array of simulated state values.
+        label (str): Label for the state (e.g., "Angle" or "Velocity").
+        color (str): Color for the simulated data plot.
+        alpha (float): Alpha value for the simulated data plot.
+    """
+    ax.plot(timespan, actual, label=f"Actual {label}", color="black", linestyle="dashed", linewidth=2)
+    ax.plot(timespan, simulated, alpha=alpha, color=color, label=f"Simulated {label}")
+    ax.set_ylabel(f"{label} (rad{'/' if label == 'Velocity' else ''}s)")
+    ax.grid(color="black", linestyle="--", linewidth=1.0, alpha=0.4)
+    ax.legend()
+
+def plot_phase_portrait(ax: plt.Axes, angle: np.ndarray, velocity: np.ndarray, simulated_angle: np.ndarray, simulated_velocity: np.ndarray, color: str = 'blue', alpha: float = 0.5) -> None:
+    """
+    Plot the phase portrait of actual and simulated data.
+
+    Args:
+        ax (plt.Axes): The matplotlib axes to plot on.
+        angle (np.ndarray): Array of actual angle values.
+        velocity (np.ndarray): Array of actual velocity values.
+        simulated_angle (np.ndarray): Array of simulated angle values.
+        simulated_velocity (np.ndarray): Array of simulated velocity values.
+        color (str): Color for the simulated data plot.
+        alpha (float): Alpha value for the simulated data plot.
+    """
+    ax.plot(angle, velocity, label="Actual", color="black", linestyle="dashed", linewidth=2)
+    ax.plot(simulated_angle, simulated_velocity, alpha=alpha, color=color, label="Simulated")
+    ax.set_xlabel("Angle (rad)")
+    ax.set_ylabel("Angular Velocity (rad/s)")
+    ax.set_title("Phase Portrait")
+    ax.grid(color="black", linestyle="--", linewidth=1.0, alpha=0.4)
+    ax.legend()
+
+def plot_simulation_errors(timespan: np.ndarray, angle: np.ndarray, velocity: np.ndarray, batched_states_trajectories: np.ndarray, predicted_terminal_points: np.ndarray, interval_terminal_states: np.ndarray, HORIZON: int, save_path: str = None, show: bool = False, title: str = "Simulation Errors", iteration: int = None) -> np.ndarray:
+    """
+    Plot simulation errors for the pendulum system and return the frame as an image.
+
+    Args:
+        timespan (np.ndarray): Array of time points.
+        angle (np.ndarray): Array of actual angle values.
+        velocity (np.ndarray): Array of actual velocity values.
+        batched_states_trajectories (np.ndarray): Array of simulated state trajectories.
+        predicted_terminal_points (np.ndarray): Array of predicted terminal points.
+        interval_terminal_states (np.ndarray): Array of actual terminal states at intervals.
+        HORIZON (int): Number of time steps in each interval.
+        save_path (str): Path to save the plot. If None, the plot is not saved.
+        show (bool): Whether to display the plot.
+        title (str): Title for the plot.
+        iteration (int, optional): Current iteration number for animation frames.
+
+    Returns:
+        np.ndarray: Image array representing the current frame.
+    """
+    fig = plt.figure(figsize=(12, 6))
+    gs = fig.add_gridspec(2, 2)
+
+    ax1 = fig.add_subplot(gs[0, 0])
+    ax2 = fig.add_subplot(gs[1, 0])
+    ax3 = fig.add_subplot(gs[:, 1])
+
+    plot_state(ax1, timespan, angle, batched_states_trajectories[:, 0], "Angle")
+    ax1.plot(timespan[HORIZON + 1 :][::HORIZON], predicted_terminal_points[:-1, 0], "ob", label="Predicted")
+    ax1.plot(timespan[HORIZON + 1 :][::HORIZON], interval_terminal_states[:, 0], "or", label="Actual")
+    if iteration is not None:
+        ax1.set_title(f"{title} (Iteration {iteration})")
+    else:
+        ax1.set_title(title)
+    ax1.legend(loc='upper right')
+
+    plot_state(ax2, timespan, velocity, batched_states_trajectories[:, 1], "Velocity")
+    ax2.plot(timespan[HORIZON + 1 :][::HORIZON], predicted_terminal_points[:-1, 1], "ob", label="Predicted")
+    ax2.plot(timespan[HORIZON + 1 :][::HORIZON], interval_terminal_states[:, 1], "or", label="Actual")
+    ax2.set_xlabel("Time (s)")
+    ax2.legend(loc='upper right')
+
+    plot_phase_portrait(ax3, angle, velocity, batched_states_trajectories[:, 0], batched_states_trajectories[:, 1])
+    ax3.plot(predicted_terminal_points[:-1, 0], predicted_terminal_points[:-1, 1], "ob", label="Predicted")
+    ax3.plot(interval_terminal_states[:, 0], interval_terminal_states[:, 1], "or", label="Actual")
+    ax3.legend(loc='upper right')
+
+    plt.tight_layout()
+    
+    if save_path:
+        os.makedirs(os.path.dirname(save_path), exist_ok=True)
+        plt.savefig(save_path, dpi=300)
+    
+    if show:
+        plt.show()
+    
+    # Convert plot to image array
+    fig.canvas.draw()
+    image = np.frombuffer(fig.canvas.tostring_rgb(), dtype='uint8')
+    image = image.reshape(fig.canvas.get_width_height()[::-1] + (3,))
+    
+    plt.close(fig)
+    
+    return image
+
+def create_animation_frame(timespan: np.ndarray, true_trajectory: np.ndarray, current_rollout: np.ndarray, iteration: int) -> np.ndarray:
+    """
+    Create a single frame for the animation of the learning process.
+
+    Args:
+        timespan (np.ndarray): Array of time points.
+        true_trajectory (np.ndarray): Array of actual state values.
+        current_rollout (np.ndarray): Array of current simulated state values.
+        iteration (int): Current iteration number.
+
+    Returns:
+        np.ndarray: Image array representing the current frame.
+    """
+    fig = plt.figure(figsize=(12, 6))  # Reduced height from 10 to 5
+    gs = fig.add_gridspec(2, 2)
+
+    ax1 = fig.add_subplot(gs[0, 0])
+    ax2 = fig.add_subplot(gs[1, 0])
+    ax3 = fig.add_subplot(gs[:, 1])
+    
+    plot_state(ax1, timespan, true_trajectory[:, 0], current_rollout[:, 0], "Angle", color="red")
+    ax1.set_title(f"Iteration {iteration}")
+
+    plot_state(ax2, timespan, true_trajectory[:, 1], current_rollout[:, 1], "Velocity", color="red")
+    ax2.set_xlabel("Time (s)")
+
+    plot_phase_portrait(ax3, true_trajectory[:, 0], true_trajectory[:, 1], current_rollout[:, 0], current_rollout[:, 1], color="red")
+    ax3.set_title("Phase Portrait")
+
+    plt.tight_layout()
+    
+    fig.canvas.draw()
+    image = np.frombuffer(fig.canvas.tostring_rgb(), dtype='uint8')
+    image = image.reshape(fig.canvas.get_width_height()[::-1] + (3,))
+    
+    plt.close(fig)
+    
+    return image
+
+def plot_full_simulation(timespan: np.ndarray, angle: np.ndarray, velocity: np.ndarray, old_rollout: np.ndarray, new_rollout: np.ndarray, save_path: str = "plots/learning_results.png", show: bool = True) -> None:
+    """
+    Plot full simulation results for the pendulum system.
+
+    Args:
+        timespan (np.ndarray): Array of time points.
+        angle (np.ndarray): Array of actual angle values.
+        velocity (np.ndarray): Array of actual velocity values.
+        old_rollout (np.ndarray): Array of simulated states using the old model.
+        new_rollout (np.ndarray): Array of simulated states using the new model.
+        save_path (str): Path to save the plot.
+        show (bool): Whether to display the plot.
+    """
+    fig = plt.figure(figsize=(12, 6))  # Reduced height from 10 to 5
+    gs = fig.add_gridspec(2, 2)
+
+    ax1 = fig.add_subplot(gs[0, 0])
+    ax2 = fig.add_subplot(gs[1, 0])
+    ax3 = fig.add_subplot(gs[:, 1])
+    
+    plot_state(ax1, timespan, angle, old_rollout[:, 0], "Angle", color="blue", alpha=0.3)
+    ax1.plot(timespan, new_rollout[:, 0], color="red", label="Optimized Model")
+    
+    plot_state(ax2, timespan, velocity, old_rollout[:, 1], "Velocity", color="blue", alpha=0.3)
+    ax2.plot(timespan, new_rollout[:, 1], color="red", label="Optimized Model")
+    ax2.set_xlabel("Time (s)")
+
+    plot_phase_portrait(ax3, angle, velocity, old_rollout[:, 0], old_rollout[:, 1], color="blue", alpha=0.3)
+    ax3.plot(new_rollout[:, 0], new_rollout[:, 1], color="red", label="Optimized Model")
+
+    plt.tight_layout()
+    if save_path:
+        os.makedirs(os.path.dirname(save_path), exist_ok=True)
+        plt.savefig(save_path, dpi=300)
+    if show:
+        plt.show()
+    else:
+        plt.close(fig)