import argparse import os from pathlib import Path from typing import Any, Dict, List, Tuple import d3il import gym import matplotlib import matplotlib.patches as patches import matplotlib.pyplot as plt import numpy as np from sklearn.linear_model import LinearRegression from sklearn.metrics import mean_squared_error, r2_score from datasets.sequence import SequenceDataset from obstacles import PILLARS matplotlib.use("Agg") class LinearDynamicsAnalyzer: def __init__(self, env_name: str, horizon: int = 16, max_episodes: int = 500): self.env_name = env_name self.horizon = horizon self.max_episodes = max_episodes self.output_dir = Path(f"logs/{self.env_name}/dynamics") self.output_dir.mkdir(parents=True, exist_ok=True) print(f"Output will be saved to: {self.output_dir}") print(f"Loading data from {self.env_name}...") ( self.X, self.y, self.actions, self.metadata, self.dataset, ) = self._collect_data() self.A: np.ndarray = None self.B: np.ndarray = None self.c: np.ndarray = None self.metrics: Dict[str, Any] = {} self.y_pred: np.ndarray = None self.verification_metrics: Dict[str, Any] = None def _collect_data( self, ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Dict[str, Any], SequenceDataset]: dataset = SequenceDataset( env=self.env_name, horizon=self.horizon, normalizer="LimitsNormalizer", preprocess_fns=[], max_path_length=200, max_n_episodes=self.max_episodes, termination_penalty=0, seed=0, ) buffer = dataset.fields x_list, y_list, actions_list = [], [], [] for path_idx in range(buffer.n_episodes): path_length = buffer.path_lengths[path_idx] observations = buffer.normed_observations[path_idx, :path_length] actions = buffer.normed_actions[path_idx, :path_length] for t in range(path_length - 1): x_list.append(observations[t]) y_list.append(observations[t + 1]) actions_list.append(actions[t]) x = np.array(x_list) y = np.array(y_list) actions = np.array(actions_list) metadata = { "env_name": self.env_name, "num_transitions": len(x), "obs_dim": x.shape[1], "action_dim": actions.shape[1] if len(actions) > 0 else 0, "num_episodes": buffer.n_episodes, } return x, y, actions, metadata, dataset def fit_model(self): print("Fitting linear dynamics model...") features = np.hstack([self.X, self.actions]) models = [] predictions = [] for i in range(self.y.shape[1]): model = LinearRegression() model.fit(features, self.y[:, i]) models.append(model) predictions.append(model.predict(features)) self.y_pred = np.column_stack(predictions) obs_dim = self.X.shape[1] self.A = np.array([m.coef_[:obs_dim] for m in models]) self.B = np.array([m.coef_[obs_dim:] for m in models]) self.c = np.array([m.intercept_ for m in models]) self.metrics = { "overall_mse": mean_squared_error(self.y, self.y_pred), "overall_r2": r2_score(self.y, self.y_pred), "dim_mse": [ mean_squared_error(self.y[:, i], self.y_pred[:, i]) for i in range(self.y.shape[1]) ], "dim_r2": [ r2_score(self.y[:, i], self.y_pred[:, i]) for i in range(self.y.shape[1]) ], } self.metrics["rmse"] = np.sqrt(self.metrics["overall_mse"]) def analyze_model(self): if self.A is None: print("Model not fitted yet. Call fit_model() first.") return print(f"\n=== Linear Dynamics Model Analysis for {self.env_name} ===") print(f"Dataset: {self.metadata['num_transitions']} transitions") print( f"Obs dim: {self.metadata['obs_dim']}, Action dim: {self.metadata['action_dim']}" ) print("\n=== Model Performance ===") print(f"Overall R²: {self.metrics['overall_r2']:.4f}") print(f"Overall RMSE: {self.metrics['rmse']:.4f}") print("\n=== Per-dimension Performance ===") for i, (r2, mse) in enumerate( zip(self.metrics["dim_r2"], self.metrics["dim_mse"]) ): print(f"Dimension {i}: R² = {r2:.4f}, MSE = {mse:.4f}") print("\n=== State Transition Matrix A (s_{t+1} dependence on s_t) ===") print("Shape:", self.A.shape) eigenvals = np.linalg.eigvals(self.A) max_eig = np.max(np.abs(eigenvals)) print(f"Spectral radius (max eigenvalue magnitude): {max_eig:.4f}") if max_eig < 1.0: print("System appears stable (spectral radius < 1.0)") else: print("System may be unstable (spectral radius >= 1.0)") print("\n=== Action Influence Matrix B (s_{t+1} dependence on a_t) ===") print("Shape:", self.B.shape) print(f"Frobenius norm: {np.linalg.norm(self.B, 'fro'):.4f}") def verify_with_env(self, n_test_cases: int = 1000): print("\n=== Verifying Linear Dynamics with Real Environment (Single-Step) ===") env = gym.make(self.env_name) env.set_seed(42) env.start() test_indices = np.random.choice(len(self.X), n_test_cases, replace=True) env_next_states, model_next_states = [], [] for test_idx in test_indices: try: current_state_norm = self.X[test_idx] action_norm = self.actions[test_idx] current_state_unnorm = self.dataset.normalizer.unnormalize( current_state_norm.reshape(1, -1), "observations" ).flatten() action_unnorm = self.dataset.normalizer.unnormalize( action_norm.reshape(1, -1), "actions" ).flatten() env.reset() env.set_state(current_state_unnorm) if not np.allclose( env.get_observation(), current_state_unnorm, atol=1e-2 ): continue next_obs_unnorm, _, _, _ = env.step(action_unnorm) next_state_model_norm = ( self.A @ current_state_norm + self.B @ action_norm + self.c ) next_state_env_norm = self.dataset.normalizer.normalize( next_obs_unnorm.reshape(1, -1), "observations" ).flatten() env_next_states.append(next_state_env_norm) model_next_states.append(next_state_model_norm) except Exception: continue if not model_next_states: print("Warning: No successful verification tests completed!") return env_next = np.array(env_next_states) model_next = np.array(model_next_states) errors = env_next - model_next self.verification_metrics = { "overall_rmse": np.sqrt(np.mean(errors**2)), "rmse_per_dim": np.sqrt(np.mean(errors**2, axis=0)), "overall_r2": r2_score(env_next, model_next), "r2_per_dim": [ r2_score(env_next[:, i], model_next[:, i]) for i in range(env_next.shape[1]) ], "n_comparisons": len(env_next), } print(f"Successful comparisons: {len(env_next)}/{n_test_cases}") print( f"Overall R² (normalized space): {self.verification_metrics['overall_r2']:.4f}" ) print( f"Overall RMSE (normalized space): {self.verification_metrics['overall_rmse']:.6f}" ) self.visualize_verification_results(env_next, model_next) def visualize(self): print("\nGenerating visualizations...") self._visualize_all_trajectories() self._visualize_model_predictions() def _visualize_all_trajectories(self): fig, ax = plt.subplots(figsize=(9, 10)) buffer = self.dataset.fields for path_idx in range(buffer.n_episodes): path_length = buffer.path_lengths[path_idx] obs_norm = buffer.normed_observations[path_idx, :path_length] obs_unnorm = self.dataset.normalizer.unnormalize(obs_norm, "observations") ax.plot(obs_unnorm[:, 2], obs_unnorm[:, 3], alpha=0.6, lw=1) ax.plot(obs_unnorm[0, 2], obs_unnorm[0, 3], "o", markersize=3, alpha=0.7) ax.plot([0.2, 0.8], [0.35, 0.35], color=[0.4, 1, 0.4], lw=5, label="Goal") for pillar in PILLARS: ax.add_patch(patches.Circle(list(pillar.center), pillar.radius, color="r", alpha=0.7)) ax.set_xlim([0.2, 0.8]) ax.set_ylim([-0.3, 0.4]) ax.set_title("All Collected Trajectories") ax.set_xlabel("X Position") ax.set_ylabel("Y Position") ax.grid(True, alpha=0.3) ax.legend() save_path = self.output_dir / "all_trajectories.png" plt.savefig(save_path, dpi=150, bbox_inches="tight") plt.close(fig) print(f"All trajectories visualization saved to: {save_path}") def _visualize_model_predictions(self): if self.y_pred is None: return obs_dim = self.y.shape[1] n_samples = min(10000, len(self.y)) idx = np.random.choice(len(self.y), n_samples, replace=False) fig, axes = plt.subplots(2, (obs_dim + 1) // 2, figsize=(12, 8), squeeze=False) axes = axes.flatten() for i in range(obs_dim): ax = axes[i] ax.scatter(self.y[idx, i], self.y_pred[idx, i], alpha=0.6, s=1) min_val, max_val = self.y[idx, i].min(), self.y[idx, i].max() ax.plot([min_val, max_val], [min_val, max_val], "r--", lw=2) ax.set_xlabel(f"True Next Obs Dim {i}") ax.set_ylabel(f"Predicted Next Obs Dim {i}") ax.set_title( f"Dim {i}: R² = {r2_score(self.y[:, i], self.y_pred[:, i]):.3f}" ) ax.grid(True, alpha=0.3) for i in range(obs_dim, len(axes)): axes[i].set_visible(False) plt.suptitle(f"Linear Model Predictions vs True Values\n{self.env_name}") plt.tight_layout() save_path = self.output_dir / "linear_model_predictions.png" plt.savefig(save_path, dpi=150, bbox_inches="tight") plt.close(fig) print(f"Prediction visualization saved to: {save_path}") def visualize_verification_results(self, env_states_all, model_states_all): if len(env_states_all) == 0 or self.verification_metrics is None: print("No verification data to visualize.") return obs_dim = env_states_all.shape[1] fig, axes = plt.subplots(2, 2, figsize=(12, 10), squeeze=False) axes = axes.flatten() for dim in range(min(obs_dim, 4)): ax = axes[dim] ax.scatter(env_states_all[:, dim], model_states_all[:, dim], alpha=0.6, s=1) min_val = min(env_states_all[:, dim].min(), model_states_all[:, dim].min()) max_val = max(env_states_all[:, dim].max(), model_states_all[:, dim].max()) ax.plot([min_val, max_val], [min_val, max_val], "r--", lw=2) rmse = self.verification_metrics["rmse_per_dim"][dim] r2 = self.verification_metrics["r2_per_dim"][dim] ax.set_xlabel(f"Environment State Dim {dim}") ax.set_ylabel(f"Model Predicted Dim {dim}") ax.set_title(f"Dim {dim}: R²={r2:.3f}, RMSE={rmse:.4f}") ax.grid(True, alpha=0.3) for dim in range(obs_dim, 4): axes[dim].set_visible(False) plt.suptitle( f"Linear Model vs Environment Verification\n" f"Overall R²: {self.verification_metrics['overall_r2']:.4f}, " f"Overall RMSE: {self.verification_metrics['overall_rmse']:.4f} " f"({self.verification_metrics['n_comparisons']} comparisons)" ) plt.tight_layout() save_path = self.output_dir / "dynamics_verification.png" plt.savefig(save_path, dpi=150, bbox_inches="tight") plt.close(fig) print(f"Verification visualization saved to: {save_path}") def save_model(self): if self.A is None: print("Model not fitted yet. Cannot save.") return save_data = { "A": self.A, "B": self.B, "c": self.c, "metadata": self.metadata, "metrics": self.metrics, "normalizer": self.dataset.normalizer, } if self.verification_metrics is not None: save_data["verification_metrics"] = self.verification_metrics model_path = self.output_dir / "linear_model.npz" np.savez(model_path, **save_data) print(f"Model saved to: {model_path}") def run_analysis(self): self.visualize() self.fit_model() self.analyze_model() self._visualize_model_predictions() self.verify_with_env() self.save_model() print(f"\n=== Summary ===") print(f"Fitted linear dynamics model for {self.env_name}") print(f"Model form: s_{{t+1}} = A*s_t + B*a_t + c") print( f"Overall performance: R² = {self.metrics['overall_r2']:.4f}, RMSE = {self.metrics['rmse']:.4f}" ) if self.verification_metrics is not None: print( f"Environment verification R²: {self.verification_metrics['overall_r2']:.4f}, " f"RMSE: {self.verification_metrics['overall_rmse']:.6f}" ) print(f"Files saved to: {self.output_dir}/") def main(): parser = argparse.ArgumentParser( description="Fit, analyze, and verify a linear dynamics model for a gym environment." ) parser.add_argument("--env", default="avoiding-v0", help="Environment name") args = parser.parse_args() analyzer = LinearDynamicsAnalyzer(env_name=args.env) analyzer.run_analysis() if __name__ == "__main__": main()