Reduced line length to 120

donk/costs/cost_function_symbolic.py

@@ -19,8 +19,8 @@ def __init__(self, cost_fun: Callable[[np.ndarray, np.ndarray], np.ndarray], T:
         """Initialize this SymbolicCostFunction.
 
         Args:
-            cost_fun: The primitive cost function. A callable which evaluates the costs for a trajectory at each timestep.
-                      Returns an ndarray with shape (T, )
+            cost_fun: The primitive cost function. A callable which evaluates the costs for a trajectory at each
+                    timestep. Returns an ndarray with shape (T, )
             T: Time horizon
             dX: Dimension of state space
             dU: Dimension of action space
@@ -103,7 +103,12 @@ class MultipartSymbolicCostFunction(SymbolicCostFunction):
     """SymbolicCostFunction composed of named partial cost functions."""
 
     def __init__(
-        self, cost_funs: list[Callable[[np.ndarray, np.ndarray], np.ndarray]], cost_function_names: list[str], T: int, dX: int, dU: int
+        self,
+        cost_funs: list[Callable[[np.ndarray, np.ndarray], np.ndarray]],
+        cost_function_names: list[str],
+        T: int,
+        dX: int,
+        dU: int,
     ) -> None:
         """Initialize this `MultipartSymbolicCostFunction`."""
         from sympy import lambdify, symbols
@@ -122,7 +127,9 @@ def cost_fun(X, U):
         U_sym = np.array(symbols(f"u:{T*dU}")).reshape(T, dU)
 
         # Lambdify and vectorize cost functions
-        self.cost_funs = [_vectorize_cost_function(lambdify([X_sym, U_sym], list(cf(X_sym, U_sym)))) for cf in cost_funs]
+        self.cost_funs = [
+            _vectorize_cost_function(lambdify([X_sym, U_sym], list(cf(X_sym, U_sym)))) for cf in cost_funs
+        ]
 
         self.cost_function_names = cost_function_names
 

donk/dynamics/linear_dynamics.py

@@ -63,7 +63,9 @@ def log_prob(self, X: np.ndarray, U: np.ndarray) -> float:
         means = self.predict(X[..., :-1, :], U, t=None)
         return multivariate_normal_logpdf(X[..., 1:, :], means, self.covar)
 
-    def evaluate(self, output_dir: str | Path, X_train: np.ndarray, U_train: np.ndarray, X_test: np.ndarray, U_test: np.ndarray):
+    def evaluate(
+        self, output_dir: str | Path, X_train: np.ndarray, U_train: np.ndarray, X_test: np.ndarray, U_test: np.ndarray
+    ):
         """Create diagnostics and evaluation plots for this dynamics model.
 
         Args:
@@ -141,7 +143,8 @@ def fit_lr(
         X: (N, T+1, dX), States
         U: (N, T, dU), Actions
         prior: DynamicsPrior to be used. May be `None` to fit without prior.
-        regularization: Eigenvalues of the joint distribution covariance are lifted to this value. Ensures matrix is not singular.
+        regularization: Eigenvalues of the joint distribution covariance are lifted to this value. Ensures matrix is
+                not singular.
         prior_weight: Weighting factor to scale the influce of the prior.
     """
     N, _, dX = X.shape

donk/models/tvlg.py

@@ -9,7 +9,9 @@ class TimeVaryingLinearGaussian:
     `p(y|t,x) ~ N(coeffcients_t * x + intercept_t, covar_t)`
     """
 
-    def __init__(self, coefficients: np.ndarray, intercept: np.ndarray, covar: np.ndarray = None, inv_covar: np.ndarray = None) -> None:
+    def __init__(
+        self, coefficients: np.ndarray, intercept: np.ndarray, covar: np.ndarray = None, inv_covar: np.ndarray = None
+    ) -> None:
         """Initialize this TimeVaryingLinearGaussian object.
 
         Must provide either covariance or precision, or both.

donk/policy/nn.py

@@ -146,7 +146,10 @@ def validation_step(model, state, action, precision):
 
         # EarlyStopping callback
         callback = tf.keras.callbacks.EarlyStopping(
-            monitor="val/loss" if X_val is not None else "train/loss", mode="min", patience=patience, restore_best_weights=True
+            monitor="val/loss" if X_val is not None else "train/loss",
+            mode="min",
+            patience=patience,
+            restore_best_weights=True,
         )
         callback.set_model(self.model)
         self.model.stop_training = False
@@ -215,7 +218,9 @@ def linearize(self, X: np.ndarray, regularization: float = 1e-6) -> LinearGaussi
             tape.watch(x)
             U = tf.reshape(self.model(x, training=False), (N, T, dU))
             U_mean = tf.reduce_mean(U, axis=0)
-            U_sum = tf.reduce_sum(U_mean, axis=0)  # Sum over timesteps as they are independend, to reduce dimension of du_dx
+            U_sum = tf.reduce_sum(
+                U_mean, axis=0
+            )  # Sum over timesteps as they are independend, to reduce dimension of du_dx
 
         # Compute and reshape jacobian
         du_dx = tape.jacobian(U_sum, x)

donk/samples/sampler.py

@@ -15,7 +15,9 @@ class Sampler(ABC):
     """
 
     @abstractmethod
-    def take_sample(self, pol: donk.policy.Policy, T: int, condition, rng: np.random.Generator = None) -> tuple[np.ndarray, np.ndarray]:
+    def take_sample(
+        self, pol: donk.policy.Policy, T: int, condition, rng: np.random.Generator = None
+    ) -> tuple[np.ndarray, np.ndarray]:
         """Take one policy sample.
 
         Args:

donk/samples/sampler_gym.py

@@ -31,7 +31,9 @@ def convert_action(self, u: np.ndarray) -> np.ndarray:
         """Convert one action from a flat numpy array to one action for the from the Gym environment."""
         return u
 
-    def take_sample(self, pol: Policy, T: int, condition: int, rng: np.random.Generator = None) -> tuple[np.ndarray, np.ndarray]:
+    def take_sample(
+        self, pol: Policy, T: int, condition: int, rng: np.random.Generator = None
+    ) -> tuple[np.ndarray, np.ndarray]:
         """Take one policy sample.
 
         Args:

donk/traj_opt/lqg.py

@@ -92,8 +92,8 @@ def optimize(self, kl_step: float, min_eta: float = 1e-6, max_eta: float = 1e16,
             kl_step: KL divergence threshold to previous policy
             min_eta: Minimal value of the Lagrangian multiplier
             max_eta: Maximal value of the Lagrangian multiplier
-            rtol: Tolerance of found solution to kl_step. Levine et al. propose a value of 0.1 in "Learning Neural Network Policies with
-                  Guided Policy Search under Unknown Dynamics", chapter 3.1
+            rtol: Tolerance of found solution to kl_step. Levine et al. propose a value of 0.1 in "Learning Neural
+                    Network Policies with Guided Policy Search under Unknown Dynamics", chapter 3.1
             full_history: Whether to return ahistory of all optimization steps, for debug purposes
 
         Returns:

donk/traj_opt/traj_opt.py

@@ -13,7 +13,11 @@ class TrajOptAlgorithm:
     """Algorithm for iterative LQR trajectory optimization."""
 
     def __init__(
-        self, kl_step: float, max_step_mult: float = 10, min_step_mult: float = 0.1, dynamics_regularization: float = 1e-6
+        self,
+        kl_step: float,
+        max_step_mult: float = 10,
+        min_step_mult: float = 0.1,
+        dynamics_regularization: float = 1e-6,
     ) -> None:
         """Initialize the iLQR trajectory optimization."""
         self.kl_step = kl_step

donk/visualization/trajectories.py

@@ -5,7 +5,9 @@
 from matplotlib.ticker import MaxNLocator
 
 
-def plot_line_set(x: np.ndarray, ys: np.ndarray, color=None, individual_lines: str | bool = "auto", label: str | None = None):
+def plot_line_set(
+    x: np.ndarray, ys: np.ndarray, color=None, individual_lines: str | bool = "auto", label: str | None = None
+):
     """Plot a set of lines with a mean and confidence interval."""
     N, _ = ys.shape
 

pyproject.toml

@@ -1,2 +1,2 @@
 [tool.black]
-line-length = 140
+line-length = 120

setup.cfg

@@ -2,9 +2,9 @@
 ignore = E203,E741,W503,D100,D104,D105,PL123,ANN101
 per-file-ignores =
     tests/*:D101,D102,ANN201
-max-line-length = 140
+max-line-length = 120
 docstring-convention = google
 
 [isort]
-line_length = 140
+line_length = 120
 profile = black

tests/dynamics_test.py

@@ -84,7 +84,10 @@ def test_fit_lr(self):
         assert_array_equal(dyn_covar, np.swapaxes(dyn_covar, 1, 2), "dyn_covar not symmetric")
         for t in range(T):
             with self.subTest(t=t):
-                self.assertTrue(all(np.linalg.eigvalsh(dyn_covar[t]) >= -1e-16), f"Negative eigenvalues {np.linalg.eigvalsh(dyn_covar[t])}")
+                self.assertTrue(
+                    all(np.linalg.eigvalsh(dyn_covar[t]) >= -1e-16),
+                    f"Negative eigenvalues {np.linalg.eigvalsh(dyn_covar[t])}",
+                )
 
     def test_fit_lr_error(self):
         from donk.dynamics.linear_dynamics import fit_lr
@@ -169,7 +172,10 @@ def test_fit_lr_with_prior(self):
         assert_array_equal(dyn_covar, np.swapaxes(dyn_covar, 1, 2), "dyn_covar not symmetric")
         for t in range(T):
             with self.subTest(t=t):
-                self.assertTrue(all(np.linalg.eigvalsh(dyn_covar[t]) >= 0), f"Negative eigenvalues {np.linalg.eigvalsh(dyn_covar[t])}")
+                self.assertTrue(
+                    all(np.linalg.eigvalsh(dyn_covar[t]) >= 0),
+                    f"Negative eigenvalues {np.linalg.eigvalsh(dyn_covar[t])}",
+                )
 
     def test_log_prob(self):
         from donk.dynamics import to_transitions

tests/samples_test.py

@@ -20,11 +20,14 @@ def test_add(self):
         pool.add(X, U)
 
         # Retrieve all
-        assert_array_equal(pool.get_transitions(), np.c_[X[:, :-1].reshape(-1, dX), U.reshape(-1, dU), X[:, 1:].reshape(-1, dX)])
+        assert_array_equal(
+            pool.get_transitions(), np.c_[X[:, :-1].reshape(-1, dX), U.reshape(-1, dU), X[:, 1:].reshape(-1, dX)]
+        )
 
         # Retrieve N
         assert_array_equal(
-            pool.get_transitions(N=5), np.c_[X[2, -6:-1].reshape(-1, dX), U[2, -5:].reshape(-1, dU), X[2, -5:].reshape(-1, dX)]
+            pool.get_transitions(N=5),
+            np.c_[X[2, -6:-1].reshape(-1, dX), U[2, -5:].reshape(-1, dU), X[2, -5:].reshape(-1, dX)],
         )