diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7f0061a..89ca86f 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,11 @@
 # Changelog
 
+## 0.6.3
+- Dashboard now synchronously updates traces of all plots when changing one plot
+- Fix some smaller issues with lag structure in data driven mpc
+- Add ADMM dashboard
+
+
 ## 0.6.2
 - Add moving horizon estimator. Use it as ``"agentlib_mpc.mhe"``
 
diff --git a/agentlib_mpc/__init__.py b/agentlib_mpc/__init__.py
index 603c022..2104f30 100644
--- a/agentlib_mpc/__init__.py
+++ b/agentlib_mpc/__init__.py
@@ -4,4 +4,4 @@
 from .modules import MODULE_TYPES
 from .models import MODEL_TYPES
 
-__version__ = "0.6.2"
+__version__ = "0.6.3"
diff --git a/agentlib_mpc/data_structures/casadi_utils.py b/agentlib_mpc/data_structures/casadi_utils.py
index daf75eb..73d0835 100644
--- a/agentlib_mpc/data_structures/casadi_utils.py
+++ b/agentlib_mpc/data_structures/casadi_utils.py
@@ -56,6 +56,8 @@ class Solvers(str, Enum):
     gurobi = "gurobi"
     bonmin = "bonmin"
     fatrop = "fatrop"
+    proxqp = "proxqp"
+    osqp = "osqp"
 
 
 class Integrators(str, Enum):
@@ -141,6 +143,10 @@ def create_solver(
             return self._create_sqpmethod_solver(nlp=nlp, options=options)
         if solver_name == Solvers.qpoases:
             return self._create_qpoases_solver(nlp=nlp, options=options)
+        if solver_name == Solvers.proxqp:
+            return self._create_proxqp_solver(nlp=nlp, options=options)
+        if solver_name == Solvers.osqp:
+            return self._create_osqp_solver(nlp=nlp, options=options)
         if solver_name == Solvers.gurobi:
             return self._create_gurobi_solver(
                 nlp=nlp, options=options, discrete=discrete
@@ -235,6 +241,26 @@ def _create_qpoases_solver(self, nlp: dict, options: dict):
         opts = {**default_opts, **options}
         return ca.qpsol("mpc", "qpoases", nlp, opts)
 
+    def _create_proxqp_solver(self, nlp: dict, options: dict):
+        default_opts = {
+            "verbose": False,
+            "print_time": False,
+            "record_time": True,
+            "proxqp": {"max_iter": 200, "eps_abs": 1e-4, "backend": "sparse"},
+        }
+        opts = {**default_opts, **options}
+        return ca.qpsol("mpc", "proxqp", nlp, opts)
+
+    def _create_osqp_solver(self, nlp: dict, options: dict):
+        default_opts = {
+            "verbose": False,
+            "print_time": False,
+            "record_time": True,
+            "osqp": {"max_iter": 200, "eps_abs": 1e-4, "verbose": False},
+        }
+        opts = {**default_opts, **options}
+        return ca.qpsol("mpc", "osqp", nlp, opts)
+
     def _create_gurobi_solver(
         self, nlp: dict, options: dict, discrete: DiscreteVars = None
     ):
diff --git a/agentlib_mpc/models/casadi_ml_model.py b/agentlib_mpc/models/casadi_ml_model.py
index cc571d6..81ef136 100644
--- a/agentlib_mpc/models/casadi_ml_model.py
+++ b/agentlib_mpc/models/casadi_ml_model.py
@@ -183,11 +183,6 @@ def __init__(self, **kwargs):
         # construct a stage function for optimization and simulation
         self.sim_step = self._make_unified_predict_function()
 
-    def specify_ann_constructed_features(self):
-        """to be implemented by user
-        todo example"""
-        pass
-
     def setup_system(self):
         return 0
 
@@ -380,9 +375,9 @@ def register_ml_models(
         for output in self.config.outputs + self.config.states:
             for serialized_output_names, ml_model in ml_model_sources_dict.items():
                 if output.name in serialized_output_names:
-                    output_to_ml_model[output.name] = (
-                        CasadiPredictor.from_serialized_model(ml_model)
-                    )
+                    output_to_ml_model[
+                        output.name
+                    ] = CasadiPredictor.from_serialized_model(ml_model)
                     ml_model_dict[output.name] = ml_model
         casadi_ml_model_dict: Dict[str, CasadiPredictor] = output_to_ml_model
         return ml_model_dict, casadi_ml_model_dict
diff --git a/agentlib_mpc/models/casadi_model.py b/agentlib_mpc/models/casadi_model.py
index c2732fc..6ac7048 100644
--- a/agentlib_mpc/models/casadi_model.py
+++ b/agentlib_mpc/models/casadi_model.py
@@ -223,7 +223,7 @@ def __attrs_post_init__(self):
     @property
     def alg(self) -> CasadiTypes:
         raise AttributeError(
-            "Casadi States should not have .alg assignments. If you wish to provide "
+            "Casadi Inputs should not have .alg assignments. If you wish to provide "
             "algebraic relationships to states, add them in the constraints."
         )
         return -1
diff --git a/agentlib_mpc/modules/dmpc/admm/admm_coordinated.py b/agentlib_mpc/modules/dmpc/admm/admm_coordinated.py
index 18cf589..3cb5999 100644
--- a/agentlib_mpc/modules/dmpc/admm/admm_coordinated.py
+++ b/agentlib_mpc/modules/dmpc/admm/admm_coordinated.py
@@ -2,8 +2,9 @@
 with a coordinator."""
 
 from collections import namedtuple
-from typing import Dict, Optional
+from typing import Dict, Optional, List
 import pandas as pd
+import pydantic
 
 from agentlib_mpc.data_structures.mpc_datamodels import MPCVariable
 from .admm import ADMM, ADMMConfig
@@ -22,6 +23,18 @@ class CoordinatedADMMConfig(MiniEmployeeConfig, ADMMConfig):
         "shared_variable_fields"
     ) + ADMMConfig.default("shared_variable_fields")
 
+    @pydantic.field_validator("couplings", "exchange")
+    def couplings_should_have_values(cls, value: List[AgentVariable]):
+        """Asserts that couplings and exchange have values, as they are needed for
+        initial guess."""
+        for var in value:
+            if var.value is None:
+                raise ValueError(
+                    "Couplings and Exchange Variables should have a value, as it is "
+                    "required for the initial guess."
+                )
+        return value
+
 
 class CoordinatedADMM(MiniEmployee, ADMM):
     """
diff --git a/agentlib_mpc/modules/mpc.py b/agentlib_mpc/modules/mpc.py
index c0d210e..2d68e77 100644
--- a/agentlib_mpc/modules/mpc.py
+++ b/agentlib_mpc/modules/mpc.py
@@ -279,9 +279,9 @@ def _init_optimization(self):
 
     def re_init_optimization(self, parameter: AgentVariable):
         """Re-initializes the optimization backend with new parameters."""
-        self.optimization_backend.discretization_options[parameter.name] = (
-            parameter.value
-        )
+        self.optimization_backend.discretization_options[
+            parameter.name
+        ] = parameter.value
         self._init_optimization()
 
     @property
diff --git a/agentlib_mpc/modules/mpc_full.py b/agentlib_mpc/modules/mpc_full.py
index 7ba7303..07cf53f 100644
--- a/agentlib_mpc/modules/mpc_full.py
+++ b/agentlib_mpc/modules/mpc_full.py
@@ -62,7 +62,7 @@ def _init_optimization(self):
             history[v] = {}
             # store scalar values as initial if they exist
             if isinstance(var.value, (float, int)):
-                timestamp = var.timestamp or self.env.now
+                timestamp = var.timestamp or self.env.time
                 value = var.value
             elif var.value is None:
                 self.logger.info(
@@ -93,7 +93,7 @@ def _remove_old_values_from_history(self):
 
             # iterate over all saved values and delete them, if they are too old
             for timestamp in list(var_history):
-                if timestamp < (self.env.now - lag_in_seconds):
+                if timestamp < (self.env.time - lag_in_seconds):
                     var_history.pop(timestamp)
 
     def _callback_hist_vars(self, variable: AgentVariable, name: str):
diff --git a/agentlib_mpc/optimization_backends/__init__.py b/agentlib_mpc/optimization_backends/__init__.py
index 5bf6cfa..469f596 100644
--- a/agentlib_mpc/optimization_backends/__init__.py
+++ b/agentlib_mpc/optimization_backends/__init__.py
@@ -41,12 +41,20 @@ def __call__(self, *args, **kwargs):
         import_path="agentlib_mpc.optimization_backends.casadi_.minlp_cia",
         class_name="CasADiCIABackend",
     ),
+    "casadi_ml": BackendImport(
+        import_path="agentlib_mpc.optimization_backends.casadi_.casadi_ml",
+        class_name="CasADiBBBackend",
+    ),
+    "casadi_admm_ml": BackendImport(
+        import_path="agentlib_mpc.optimization_backends.casadi_.casadi_admm_ml",
+        class_name="CasADiADMMBackend_NN",
+    ),
     "casadi_nn": BackendImport(
-        import_path="agentlib_mpc.optimization_backends.casadi_.casadi_nn",
+        import_path="agentlib_mpc.optimization_backends.casadi_.casadi_ml",
         class_name="CasADiBBBackend",
     ),
     "casadi_admm_nn": BackendImport(
-        import_path="agentlib_mpc.optimization_backends.casadi_.casadi_admm_nn",
+        import_path="agentlib_mpc.optimization_backends.casadi_.casadi_admm_ml",
         class_name="CasADiADMMBackend_NN",
     ),
     "casadi_mhe": BackendImport(
diff --git a/agentlib_mpc/optimization_backends/casadi_/admm.py b/agentlib_mpc/optimization_backends/casadi_/admm.py
index 53f5a97..068c7d0 100644
--- a/agentlib_mpc/optimization_backends/casadi_/admm.py
+++ b/agentlib_mpc/optimization_backends/casadi_/admm.py
@@ -257,7 +257,7 @@ def _discretize(self, sys: CasadiADMMSystem):
             # integral and multiple shooting constraint
             fk = opt_integrator(
                 x0=xk,
-                p=ca.vertcat(uk, v_localk, dk, const_par),
+                p=ca.vertcat(uk, v_localk, dk, const_par, zk, yk),
             )
             xk_end = fk["xf"]
             self.k += 1
@@ -289,10 +289,17 @@ def _create_ode(
             sys.local_couplings.full_symbolic,
             sys.non_controlled_inputs.full_symbolic,
             sys.model_parameters.full_symbolic,
+            sys.algebraics.full_symbolic,
+            sys.outputs.full_symbolic,
         )
         integrator_ode = {"x": x, "p": p, "ode": ode}
-        opt_integrator = ca.integrator("system", integrator, integrator_ode, opts)
-
+        if integrator == Integrators.euler:
+            xk_end = x + ode * opts["tf"]
+            opt_integrator = ca.Function(
+                "system", [x, p], [xk_end], ["x0", "p"], ["xf"]
+            )
+        else:  # rk, cvodes
+            opt_integrator = ca.integrator("system", integrator, integrator_ode, opts)
         return opt_integrator
 
 
diff --git a/agentlib_mpc/optimization_backends/casadi_/basic.py b/agentlib_mpc/optimization_backends/casadi_/basic.py
index fc9fd78..8710395 100644
--- a/agentlib_mpc/optimization_backends/casadi_/basic.py
+++ b/agentlib_mpc/optimization_backends/casadi_/basic.py
@@ -430,7 +430,7 @@ def _discretize(self, sys: BaseSystem):
             )
             fk = opt_integrator(
                 x0=xk,
-                p=ca.vertcat(uk, dk, const_par),
+                p=ca.vertcat(uk, dk, const_par, zk, yk),
             )
             xk_end = fk["xf"]
             # calculate model constraint
@@ -449,10 +449,13 @@ def _create_ode(self, sys: BaseSystem, opts: dict, integrator: Integrators):
         ode = sys.ode
         # create inputs
         x = sys.states.full_symbolic
+        # the order of elements here is important when calling the integrator!
         p = ca.vertcat(
             sys.controls.full_symbolic,
             sys.non_controlled_inputs.full_symbolic,
             sys.model_parameters.full_symbolic,
+            sys.algebraics.full_symbolic,
+            sys.outputs.full_symbolic,
         )
         integrator_ode = {"x": x, "p": p, "ode": ode}
 
diff --git a/agentlib_mpc/optimization_backends/casadi_/casadi_admm_nn.py b/agentlib_mpc/optimization_backends/casadi_/casadi_admm_ml.py
similarity index 98%
rename from agentlib_mpc/optimization_backends/casadi_/casadi_admm_nn.py
rename to agentlib_mpc/optimization_backends/casadi_/casadi_admm_ml.py
index ab8e76d..3ee18f3 100644
--- a/agentlib_mpc/optimization_backends/casadi_/casadi_admm_nn.py
+++ b/agentlib_mpc/optimization_backends/casadi_/casadi_admm_ml.py
@@ -12,10 +12,10 @@
 )
 from agentlib_mpc.data_structures.ml_model_datatypes import name_with_lag
 from agentlib_mpc.models.casadi_ml_model import CasadiMLModel
-from agentlib_mpc.optimization_backends.casadi_.casadi_nn import (
-    CasadiNNSystem,
+from agentlib_mpc.optimization_backends.casadi_.casadi_ml import (
+    CasadiMLSystem,
     CasADiBBBackend,
-    MultipleShooting_NN,
+    MultipleShooting_ML,
 )
 from agentlib_mpc.optimization_backends.casadi_.core.VariableGroup import (
     OptimizationVariable,
@@ -31,7 +31,7 @@
 logger = logging.getLogger(__name__)
 
 
-class CasadiADMMNNSystem(CasadiADMMSystem, CasadiNNSystem):
+class CasadiADMMNNSystem(CasadiADMMSystem, CasadiMLSystem):
     """
     In this class, the lags are determined by the trainer alone and the lags are
     saved in the serialized MLModel so that it doesn't have to be defined in the
@@ -236,7 +236,7 @@ def sim_step_quantities(
         }
 
 
-class MultipleShootingADMMNN(ADMMMultipleShooting, MultipleShooting_NN):
+class MultipleShootingADMMNN(ADMMMultipleShooting, MultipleShooting_ML):
     max_lag: int
 
     def _discretize(self, sys: CasadiADMMNNSystem):
@@ -331,9 +331,6 @@ def _discretize(self, sys: CasadiADMMNNSystem):
         self.pred_time += ts
         mx_dict[self.pred_time] = {sys.states.name: self.add_opt_var(sys.states)}
 
-        # control of last time step
-        mx_dict[0 - ts][sys.controls.name] = self.add_opt_par(sys.last_control)
-
         all_quantities = sys.all_system_quantities()
         # add constraints and create the objective function for all stages
         for time in prediction_grid:
diff --git a/agentlib_mpc/optimization_backends/casadi_/casadi_nn.py b/agentlib_mpc/optimization_backends/casadi_/casadi_ml.py
similarity index 95%
rename from agentlib_mpc/optimization_backends/casadi_/casadi_nn.py
rename to agentlib_mpc/optimization_backends/casadi_/casadi_ml.py
index 6785c9c..70a753a 100644
--- a/agentlib_mpc/optimization_backends/casadi_/casadi_nn.py
+++ b/agentlib_mpc/optimization_backends/casadi_/casadi_ml.py
@@ -28,7 +28,7 @@
 from agentlib_mpc.optimization_backends.casadi_.full import FullSystem
 
 
-class CasadiNNSystem(FullSystem):
+class CasadiMLSystem(FullSystem):
     # multiple possibilities of using the MLModel
     # stage function for neural networks
     model: CasadiMLModel
@@ -73,14 +73,14 @@ def initialize(self, model: CasadiMLModel, var_ref: FullVariableReference):
             ref_list=var_ref.parameters,
         )
         self.initial_state = OptimizationParameter.declare(
-            denotation="initial_state",  # append the 0 as a convention to get initial guess
+            denotation="initial_state",
             variables=model.get_states(var_ref.states),
             ref_list=var_ref.states,
             use_in_stage_function=False,
             assert_complete=True,
         )
         self.last_control = OptimizationParameter.declare(
-            denotation="initial_control",  # append the 0 as a convention to get initial guess
+            denotation="initial_control",
             variables=model.get_inputs(var_ref.controls),
             ref_list=var_ref.controls,
             use_in_stage_function=False,
@@ -115,10 +115,10 @@ def all_system_quantities(self) -> dict[str, OptimizationQuantity]:
         return {var.name: var for var in self.quantities}
 
 
-class MultipleShooting_NN(MultipleShooting):
+class MultipleShooting_ML(MultipleShooting):
     max_lag: int
 
-    def _discretize(self, sys: CasadiNNSystem):
+    def _discretize(self, sys: CasadiMLSystem):
         n = self.options.prediction_horizon
         ts = self.options.time_step
         const_par = self.add_opt_par(sys.model_parameters)
@@ -147,6 +147,11 @@ def _discretize(self, sys: CasadiNNSystem):
                 sys.states, lb=x_past, ub=x_past, guess=x_past
             )
             mx_dict[time][sys.initial_state.name] = x_past
+            y_past = self.add_opt_par(sys.initial_output)
+            mx_dict[time][sys.outputs.name] = self.add_opt_var(
+                sys.outputs, lb=y_past, ub=y_past, guess=y_past
+            )
+            mx_dict[time][sys.initial_output.name] = y_past
 
         # add past inputs
         for time in pre_grid_inputs:
@@ -179,9 +184,6 @@ def _discretize(self, sys: CasadiNNSystem):
         self.pred_time += ts
         mx_dict[self.pred_time] = {sys.states.name: self.add_opt_var(sys.states)}
 
-        # control of last time step
-        mx_dict[0 - ts][sys.controls.name] = self.add_opt_par(sys.last_control)
-
         all_quantities = sys.all_system_quantities()
         # add constraints and create the objective function for all stages
         for time in prediction_grid:
@@ -231,13 +233,13 @@ def _discretize(self, sys: CasadiNNSystem):
             )
             self.objective_function += stage_result["cost_function"] * ts
 
-    def initialize(self, system: CasadiNNSystem, solver_factory: SolverFactory):
+    def initialize(self, system: CasadiMLSystem, solver_factory: SolverFactory):
         """Initializes the trajectory optimization problem, creating all symbolic
         variables of the OCP, the mapping function and the numerical solver."""
         self._construct_stage_function(system)
         super().initialize(system=system, solver_factory=solver_factory)
 
-    def _construct_stage_function(self, system: CasadiNNSystem):
+    def _construct_stage_function(self, system: CasadiMLSystem):
         """
         Combine information from the model and the var_ref to create CasADi
         functions which describe the system dynamics and constraints at each
@@ -343,7 +345,7 @@ def _construct_stage_function(self, system: CasadiNNSystem):
             output_denotations,
         )
 
-    def _create_lag_structure_for_denotations(self, system: CasadiNNSystem):
+    def _create_lag_structure_for_denotations(self, system: CasadiMLSystem):
         all_system_quantities = self.all_system_quantities(system)
         all_input_variables = {}
         lagged_inputs: dict[int, dict[str, ca.MX]] = {}
@@ -381,9 +383,9 @@ class CasADiBBBackend(CasADiBaseBackend):
     Class doing optimization with a MLModel.
     """
 
-    system_type = CasadiNNSystem
-    discretization_types = {DiscretizationMethod.multiple_shooting: MultipleShooting_NN}
-    system: CasadiNNSystem
+    system_type = CasadiMLSystem
+    discretization_types = {DiscretizationMethod.multiple_shooting: MultipleShooting_ML}
+    system: CasadiMLSystem
     # a dictionary of collections of the variable lags
     lag_collection: Dict[str, collections.deque] = {}
     max_lag: int
diff --git a/agentlib_mpc/optimization_backends/casadi_/core/casadi_backend.py b/agentlib_mpc/optimization_backends/casadi_/core/casadi_backend.py
index c4c104e..b333fc8 100644
--- a/agentlib_mpc/optimization_backends/casadi_/core/casadi_backend.py
+++ b/agentlib_mpc/optimization_backends/casadi_/core/casadi_backend.py
@@ -1,10 +1,9 @@
 import logging
 import platform
 from pathlib import Path
-from typing import Callable, Type, Optional
+from typing import Type, Optional
 
 import casadi as ca
-import pandas as pd
 import pydantic
 from agentlib.core.errors import ConfigurationError
 
diff --git a/agentlib_mpc/optimization_backends/casadi_/core/discretization.py b/agentlib_mpc/optimization_backends/casadi_/core/discretization.py
index cc82706..646ed2a 100644
--- a/agentlib_mpc/optimization_backends/casadi_/core/discretization.py
+++ b/agentlib_mpc/optimization_backends/casadi_/core/discretization.py
@@ -2,7 +2,6 @@
 
 import abc
 import dataclasses
-from functools import cached_property
 from pathlib import Path
 from typing import TypeVar, Union, Callable, Optional
 
@@ -65,7 +64,7 @@ def variable_lookup(self) -> dict[str, int]:
                 lookup[label[1]] = index
         return lookup
 
-    @cached_property
+    @property
     def df(self) -> pd.DataFrame:
         return pd.DataFrame(self.matrix, index=self.grid, columns=self.columns)
 
diff --git a/agentlib_mpc/optimization_backends/casadi_/full.py b/agentlib_mpc/optimization_backends/casadi_/full.py
index 95e114e..61c1699 100644
--- a/agentlib_mpc/optimization_backends/casadi_/full.py
+++ b/agentlib_mpc/optimization_backends/casadi_/full.py
@@ -152,7 +152,7 @@ def _discretize(self, sys: FullSystem):
             # integral and multiple shooting constraint
             fk = opt_integrator(
                 x0=xk,
-                p=ca.vertcat(uk, dk, const_par),
+                p=ca.vertcat(uk, dk, const_par, zk, yk),
             )
             xk_end = fk["xf"]
             self.k += 1
diff --git a/agentlib_mpc/utils/plotting/admm_dashboard.py b/agentlib_mpc/utils/plotting/admm_dashboard.py
new file mode 100644
index 0000000..d1c8a5c
--- /dev/null
+++ b/agentlib_mpc/utils/plotting/admm_dashboard.py
@@ -0,0 +1,585 @@
+import os
+import webbrowser
+from pathlib import Path
+from typing import List, Dict, Optional, Literal
+
+from agentlib.core.errors import OptionalDependencyError
+import pandas as pd
+
+
+from agentlib_mpc.utils.analysis import load_mpc
+from agentlib_mpc.utils.plotting.admm_residuals import load_residuals
+from agentlib_mpc.utils.plotting.interactive import get_port
+
+try:
+    import dash
+    from dash import html, dcc
+    from dash.dependencies import Input, Output, State
+    import plotly.graph_objects as go
+    import dash_daq as daq
+except ImportError as e:
+    raise OptionalDependencyError(
+        dependency_name="interactive",
+        dependency_install="plotly, dash",
+        used_object="interactive",
+    ) from e
+
+
+def load_agent_data(directory: str) -> Dict[str, pd.DataFrame]:
+    """
+    Load MPC data for multiple agents from files containing 'admm' in their name.
+
+    Args:
+    directory (str): Directory path containing the data files.
+
+    Returns:
+    Dict[str, pd.DataFrame]: Dictionary with agent names as keys and their data as values.
+    """
+    agent_data = {}
+    for filename in os.listdir(directory):
+        if (
+            "admm" in filename.casefold()
+            and filename.endswith(".csv")
+            and not "stats" in filename.casefold()
+        ):
+            file_path = os.path.join(directory, filename)
+            agent_name = f"Agent_{len(agent_data) + 1}"
+            try:
+                agent_data[agent_name] = load_mpc(file_path)
+            except Exception as e:
+                print(f"Error loading file {filename}: {str(e)}")
+    return agent_data
+
+
+def get_coupling_variables(df: pd.DataFrame) -> List[str]:
+    """
+    Identify coupling variables in the dataframe.
+
+    Args:
+    df (pd.DataFrame): The MPC data for an agent.
+
+    Returns:
+    List[str]: List of coupling variable names.
+    """
+    coupling_vars = []
+    for col in df.columns:
+        if col[0] == "parameter" and col[1].startswith("admm_coupling_mean_"):
+            var_name = col[1].replace("admm_coupling_mean_", "")
+            if ("variable", var_name) in df.columns:
+                coupling_vars.append(var_name)
+    return coupling_vars
+
+
+def get_data_for_plot(
+    agent_data: Dict[str, pd.DataFrame],
+    time_step: float,
+    iteration: int,
+    coupling_var: str,
+) -> Dict[str, List[float]]:
+    """
+    Extract data for the coupling variable plot.
+
+    Args:
+    agent_data (Dict[str, pd.DataFrame]): Dictionary containing data for each agent.
+    time_step (float): Selected time step.
+    iteration (int): Selected iteration number.
+    coupling_var (str): Name of the selected coupling variable.
+
+    Returns:
+    Dict[str, List[float]]: Dictionary with agent names as keys and their values as lists.
+    """
+    plot_data = {}
+    prediction_grid = None
+
+    for agent_name, df in agent_data.items():
+        try:
+            agent_data_at_step = df.loc[(time_step, iteration)]
+            agent_values = agent_data_at_step[("variable", coupling_var)].values
+            plot_data[agent_name] = agent_values.tolist()
+
+            if prediction_grid is None:
+                prediction_grid = agent_data_at_step.index.tolist()
+
+            # Get mean value (assuming it's the same for all agents)
+            if "Mean" not in plot_data:
+                mean_values = agent_data_at_step[
+                    ("parameter", f"admm_coupling_mean_{coupling_var}")
+                ].values
+                plot_data["Mean"] = mean_values.tolist()
+        except KeyError:
+            continue  # Skip this agent if data is not available for the selected time step and iteration
+
+    return plot_data, prediction_grid
+
+
+def create_coupling_var_plot(
+    plot_data: Dict[str, List[float]], prediction_grid: List[float], coupling_var: str
+) -> go.Figure:
+    """
+    Create a plotly figure for the coupling variable plot.
+
+    Args:
+    plot_data (Dict[str, List[float]]): Dictionary with agent names as keys and their values as lists.
+    prediction_grid (List[float]): List of prediction grid values.
+    coupling_var (str): Name of the coupling variable.
+
+    Returns:
+    go.Figure: Plotly figure object.
+    """
+    fig = go.Figure()
+
+    for agent_name, values in plot_data.items():
+        if agent_name == "Mean":
+            line_style = dict(color="red", dash="dash", width=2)
+        else:
+            line_style = dict(width=1)
+
+        fig.add_trace(
+            go.Scatter(
+                x=prediction_grid,
+                y=values,
+                mode="lines",
+                name=agent_name,
+                line=line_style,
+            )
+        )
+
+    fig.update_layout(
+        title=f"Coupling Variable: {coupling_var}",
+        xaxis_title="Prediction Grid",
+        yaxis_title="Value",
+        legend_title="Legend",
+        legend=dict(yanchor="top", y=0.99, xanchor="left", x=0.01),
+    )
+
+    return fig
+
+
+def get_max_iterations_per_timestep(
+    agent_data: Dict[str, pd.DataFrame],
+) -> Dict[float, int]:
+    max_iterations = {}
+    for df in agent_data.values():
+        for time_step in df.index.get_level_values(0).unique():
+            iterations = df.loc[time_step].index.get_level_values(0).max()
+            if (
+                time_step not in max_iterations
+                or iterations > max_iterations[time_step]
+            ):
+                max_iterations[time_step] = iterations
+    return max_iterations
+
+
+def create_residuals_plot(residuals_df: pd.DataFrame, time_step: float) -> go.Figure:
+    """
+    Create a plotly figure for the residuals plot.
+
+    Args:
+    residuals_df (pd.DataFrame): DataFrame containing residuals data.
+    time_step (float): Selected time step.
+
+    Returns:
+    go.Figure: Plotly figure object.
+    """
+    fig = go.Figure()
+
+    residuals_data = residuals_df.loc[time_step]
+
+    if len(residuals_data) == 1:  # Only one iteration (iteration = 0)
+        primal_residual = residuals_data["primal_residual"].iloc[0]
+        dual_residual = residuals_data["dual_residual"].iloc[0]
+
+        fig.add_trace(
+            go.Scatter(
+                x=[0, 1],
+                y=[primal_residual, primal_residual],
+                mode="lines",
+                name="Primal Residual",
+            )
+        )
+
+        fig.add_trace(
+            go.Scatter(
+                x=[0, 1],
+                y=[dual_residual, dual_residual],
+                mode="lines",
+                name="Dual Residual",
+            )
+        )
+
+        fig.update_layout(
+            xaxis_range=[0, 1],
+            xaxis_title="Iteration",
+        )
+    else:
+        fig.add_trace(
+            go.Scatter(
+                x=residuals_data.index,
+                y=residuals_data["primal_residual"],
+                mode="lines",
+                name="Primal Residual",
+            )
+        )
+
+        fig.add_trace(
+            go.Scatter(
+                x=residuals_data.index,
+                y=residuals_data["dual_residual"],
+                mode="lines",
+                name="Dual Residual",
+            )
+        )
+
+        fig.update_layout(
+            xaxis_title="Iteration",
+        )
+
+    fig.update_layout(
+        title="Primal and Dual Residuals",
+        yaxis_title="Residual Value",
+        yaxis_type="log",
+        yaxis=dict(
+            tickformat=".2e",  # Use scientific notation with 2 decimal places
+            exponentformat="e",  # Use "e" notation for exponents
+        ),
+        legend_title="Legend",
+    )
+
+    return fig
+
+
+def create_app(agent_data: Dict[str, pd.DataFrame], residuals_df: pd.DataFrame):
+    """
+    Create and configure the Dash app.
+
+    Args:
+    agent_data (Dict[str, pd.DataFrame]): Dictionary containing data for each agent.
+    residuals_df (pd.DataFrame): DataFrame containing residuals data.
+
+    Returns:
+    dash.Dash: Configured Dash app.
+    """
+    app = dash.Dash(__name__)
+
+    # Get time steps and iteration numbers
+    first_agent_data = next(iter(agent_data.values()))
+    time_steps = sorted(first_agent_data.index.get_level_values(0).unique())
+    max_iterations_per_timestep = get_max_iterations_per_timestep(agent_data)
+    overall_max_iterations = max(max_iterations_per_timestep.values())
+
+    # Get coupling variables
+    coupling_vars = get_coupling_variables(first_agent_data)
+
+    app.layout = html.Div(
+        [
+            html.H1("Distributed MPC with ADMM Dashboard"),
+            # time step slider
+            html.Div(
+                [
+                    html.Label("Time Step"),
+                    html.Div(
+                        [
+                            html.Div(
+                                [
+                                    dcc.Slider(
+                                        id="time-step-slider",
+                                        min=0,
+                                        max=len(time_steps) - 1,
+                                        value=0,
+                                        marks={
+                                            i: f"{time_steps[i]:.2f}"
+                                            for i in range(
+                                                0,
+                                                len(time_steps),
+                                                max(1, len(time_steps) // 10),
+                                            )
+                                        },
+                                        step=1,
+                                    )
+                                ],
+                                style={
+                                    "width": "80%",
+                                    "display": "inline-block",
+                                    "verticalAlign": "middle",
+                                },
+                            ),
+                            html.Div(
+                                [
+                                    html.Button(
+                                        "◀",
+                                        id="prev-time-step",
+                                        n_clicks=0,
+                                        style={"marginRight": "5px"},
+                                    ),
+                                    html.Button(
+                                        "▶",
+                                        id="next-time-step",
+                                        n_clicks=0,
+                                        style={"marginRight": "5px"},
+                                    ),
+                                    html.Div(
+                                        id="time-step-display",
+                                        style={
+                                            "display": "inline-block",
+                                            "marginRight": "10px",
+                                        },
+                                    ),
+                                    daq.NumericInput(
+                                        id="time-step-input",
+                                        min=0,
+                                        max=len(time_steps) - 1,
+                                        value=0,
+                                        size=60,
+                                    ),
+                                ],
+                                style={
+                                    "width": "20%",
+                                    "display": "inline-block",
+                                    "verticalAlign": "middle",
+                                    "textAlign": "right",
+                                },
+                            ),
+                        ]
+                    ),
+                ]
+            ),
+            # iteration slide
+            html.Div(
+                [
+                    html.Label("Iteration Number"),
+                    html.Div(
+                        [
+                            html.Div(
+                                [
+                                    dcc.Slider(
+                                        id="iteration-slider",
+                                        min=0,
+                                        max=overall_max_iterations,
+                                        value=0,
+                                        marks={
+                                            i: str(i)
+                                            for i in range(
+                                                0,
+                                                overall_max_iterations + 1,
+                                                max(1, overall_max_iterations // 10),
+                                            )
+                                        },
+                                        step=1,
+                                    )
+                                ],
+                                style={
+                                    "width": "80%",
+                                    "display": "inline-block",
+                                    "verticalAlign": "middle",
+                                },
+                            ),
+                            html.Div(
+                                [
+                                    daq.NumericInput(
+                                        id="iteration-input",
+                                        min=0,
+                                        max=overall_max_iterations,
+                                        value=0,
+                                        size=60,
+                                    )
+                                ],
+                                style={
+                                    "width": "20%",
+                                    "display": "inline-block",
+                                    "verticalAlign": "middle",
+                                    "textAlign": "right",
+                                },
+                            ),
+                        ]
+                    ),
+                ]
+            ),
+            html.Div(
+                [
+                    html.Label("Coupling Variable"),
+                    dcc.Dropdown(
+                        id="coupling-var-dropdown",
+                        options=[{"label": var, "value": var} for var in coupling_vars],
+                        value=coupling_vars[0] if coupling_vars else None,
+                    ),
+                ]
+            ),
+            dcc.Graph(id="coupling-var-plot"),
+            dcc.Graph(id="residuals-plot"),
+            dcc.Store(id="y-axis-range"),
+        ]
+    )
+
+    @app.callback(
+        [
+            Output("time-step-input", "value"),
+            Output("time-step-display", "children"),
+            Output("time-step-slider", "value"),
+        ],
+        [
+            Input("time-step-slider", "value"),
+            Input("prev-time-step", "n_clicks"),
+            Input("next-time-step", "n_clicks"),
+            Input("time-step-input", "value"),
+        ],
+        [State("time-step-input", "value")],
+    )
+    def update_time_step(
+        slider_value, prev_clicks, next_clicks, input_value, current_value
+    ):
+        ctx = dash.callback_context
+        if not ctx.triggered:
+            return (
+                current_value,
+                f"Current time step: {time_steps[current_value]:.2f}",
+                current_value,
+            )
+
+        input_id = ctx.triggered[0]["prop_id"].split(".")[0]
+
+        if input_id == "time-step-slider":
+            new_value = slider_value
+        elif input_id == "prev-time-step":
+            new_value = max(0, current_value - 1)
+        elif input_id == "next-time-step":
+            new_value = min(len(time_steps) - 1, current_value + 1)
+        elif input_id == "time-step-input":
+            new_value = input_value
+        else:
+            new_value = current_value
+
+        return new_value, f"Current time step: {time_steps[new_value]:.2f}", new_value
+
+    @app.callback(
+        [
+            Output("iteration-slider", "max"),
+            Output("iteration-slider", "marks"),
+            Output("iteration-input", "max"),
+        ],
+        [Input("time-step-input", "value")],
+    )
+    def update_iteration_range(time_step_index):
+        time_step = time_steps[time_step_index]
+        max_iter = max_iterations_per_timestep[time_step]
+        marks = {i: str(i) for i in range(0, max_iter + 1, max(1, max_iter // 10))}
+        return max_iter, marks, max_iter
+
+    @app.callback(
+        Output("coupling-var-plot", "figure"),
+        [
+            Input("time-step-input", "value"),
+            Input("iteration-input", "value"),
+            Input("coupling-var-dropdown", "value"),
+            Input("y-axis-range", "data"),
+        ],
+    )
+    def update_coupling_var_plot(time_step_index, iteration, coupling_var, y_range):
+        if coupling_var is None:
+            return go.Figure()
+
+        time_step = time_steps[time_step_index]
+        max_iter = max_iterations_per_timestep[time_step]
+        iteration = min(iteration, max_iter)
+
+        plot_data, prediction_grid = get_data_for_plot(
+            agent_data, time_step, iteration, coupling_var
+        )
+        fig = create_coupling_var_plot(plot_data, prediction_grid, coupling_var)
+
+        if y_range is not None:
+            fig.update_layout(yaxis_range=y_range)
+
+        return fig
+
+    @app.callback(
+        Output("y-axis-range", "data"),
+        [Input("time-step-input", "value"), Input("coupling-var-dropdown", "value")],
+    )
+    def compute_y_axis_range(time_step_index, coupling_var):
+        if coupling_var is None:
+            return None
+
+        time_step = time_steps[time_step_index]
+
+        max_vals = []
+        min_vals = []
+        for agent, data in agent_data.items():
+            try:
+                step_data = data[("variable", coupling_var)][time_step]
+            except KeyError:
+                continue
+            max_vals.append(step_data.max())
+            min_vals.append(step_data.min())
+
+        y_min = min(min_vals)
+        y_max = max(max_vals)
+
+        y_range = [y_min - 0.1 * (y_max - y_min), y_max + 0.1 * (y_max - y_min)]
+        return y_range
+
+    @app.callback(
+        Output("residuals-plot", "figure"),
+        [Input("time-step-input", "value")],
+    )
+    def update_residuals_plot(time_step_index):
+        time_step = time_steps[time_step_index]
+
+        # Check if residuals data exists for this time step
+        if time_step not in residuals_df.index:
+            # If no data, return an empty figure with a message
+            fig = go.Figure()
+            fig.add_annotation(
+                text="No residuals data available for this time step",
+                xref="paper",
+                yref="paper",
+                x=0.5,
+                y=0.5,
+                showarrow=False,
+            )
+            return fig
+
+        return create_residuals_plot(residuals_df, time_step)
+
+    return app
+
+
+def main():
+    # Specify the directory containing the data files
+    data_directory = Path(r"D:\repos\juelich_mpc\juelich_mpc\mpc\simple_model\results")
+
+    # Load agent data
+    agent_data = load_agent_data(data_directory)
+    agent_data["heating"] = load_mpc(Path(data_directory, "heating_agent_res.csv"))
+    # "room_1": load_mpc(Path(data_directory, "room_1_admm.csv")),
+    # "room_2": load_mpc(Path(data_directory, "room_2_admm.csv")),
+    # "room_3": load_mpc(Path(data_directory, "room_3_admm.csv")),
+    # "heating": load_mpc(Path(data_directory, "heating_agent_res.csv")),
+
+    # Load residuals data
+    residuals_file = os.path.join(
+        data_directory, "residuals.csv"
+    )  # Adjust the filename as needed
+    residuals_df = load_residuals(residuals_file)
+
+    # Create and run the app
+    app = create_app(agent_data, residuals_df)
+    port = get_port()
+
+    webbrowser.open_new_tab(f"http://localhost:{port}")
+    app.run_server(debug=False, port=port)
+
+
+def show_admm_dashboard(
+    data: dict[str, pd.DataFrame],
+    residuals: Optional[pd.DataFrame] = None,
+    scale: Literal["seconds", "minutes", "hours", "days"] = "seconds",
+):
+    app = create_app(data, residuals)
+    port = get_port()
+
+    webbrowser.open_new_tab(f"http://localhost:{port}")
+    app.run_server(debug=False, port=port)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/agentlib_mpc/utils/plotting/interactive.py b/agentlib_mpc/utils/plotting/interactive.py
index 1156ae3..ec48255 100644
--- a/agentlib_mpc/utils/plotting/interactive.py
+++ b/agentlib_mpc/utils/plotting/interactive.py
@@ -14,7 +14,7 @@
 try:
     import dash
     from dash import html, dcc
-    from dash.dependencies import Input, Output
+    from dash.dependencies import Input, Output, State
     import plotly.graph_objects as go
 except ImportError as e:
     raise OptionalDependencyError(
@@ -173,6 +173,7 @@ def plot_mpc_plotly(
         ),
         xaxis_title=x_axis_label,
         yaxis_title=y_axis_label,
+        uirevision="same",  # Add this line
     )
 
     return fig
@@ -230,47 +231,73 @@ def show_dashboard(
         except Exception:
             pass
 
+    # Store initial figures
+    initial_figures = {}
+    for column in columns_okay:
+        fig = plot_mpc_plotly(
+            data["variable"][column],
+            convert_to=scale,
+            y_axis_label=column,
+        )
+        # Add uirevision to maintain legend state
+        fig.update_layout(uirevision="same")
+        initial_figures[column] = fig
+
     # Define the layout of the webpage
     app.layout = html.Div(
         [
             html.H1("MPC Results"),
+            # Store for keeping track of trace visibility
+            dcc.Store(id="trace-visibility", data={}),
             make_components(columns_okay, data, stats=stats, convert_to=scale),
         ]
     )
 
     port = get_port()
 
-    webbrowser.open_new_tab(f"http://localhost:{port}")
-    app.run(debug=False, port=port)
-
-    # Create a callback to update the visibility of traces in all plots
     @app.callback(
-        [Output(f"plot-{column}", "figure") for column in columns],
-        [Input(f"plot-{column}", "clickData") for column in columns],
+        [Output(f"plot-{column}", "figure") for column in columns_okay],
+        [Input(f"plot-{column}", "restyleData") for column in columns_okay],
+        [State(f"plot-{column}", "figure") for column in columns_okay],
     )
-    def update_trace_visibility(*args):
-        clickData = args[: len(columns)]
-        figures = args[len(columns) :]
+    def update_plots(*args):
+        ctx = dash.callback_context
+        if not ctx.triggered:
+            return [dash.no_update] * len(columns_okay)
+
+        n_plots = len(columns_okay)
+        restyle_data = args[:n_plots]
+        current_figures = args[n_plots:]
+
+        # Find which plot was changed
+        triggered_prop = ctx.triggered[0]["prop_id"].split(".")[0]
+        triggered_index = next(
+            i for i, col in enumerate(columns_okay) if f"plot-{col}" == triggered_prop
+        )
+        triggered_data = restyle_data[triggered_index]
 
-        print(f"in callback, {figures} and clickdata {clickData}")
+        if not triggered_data:
+            return [dash.no_update] * n_plots
 
-        # Get the name of the clicked legend item
-        clicked_legend = None
-        for click in clickData:
-            if click and "curveNumber" in click:
-                clicked_legend = figures[0]["data"][click["curveNumber"]]["name"]
-                break
+        # Get the visibility update from the triggered plot
+        visibility_update = triggered_data[0].get("visible", [None])[0]
+        trace_indices = triggered_data[1]
 
-        # Update the visibility of traces in all plots based on the clicked legend item
+        # Update all figures
         updated_figures = []
-        for fig in figures:
-            for trace in fig["data"]:
-                if trace["name"] == clicked_legend:
-                    trace["visible"] = not trace["visible"]
+        for fig in current_figures:
+            # Ensure uirevision is set
+            fig["layout"]["uirevision"] = "same"
+            # Update visibility for the corresponding traces
+            for idx in trace_indices:
+                fig["data"][idx]["visible"] = visibility_update
             updated_figures.append(fig)
 
         return updated_figures
 
+    webbrowser.open_new_tab(f"http://localhost:{port}")
+    app.run(debug=False, port=port)
+
 
 def make_components(
     columns, data, convert_to, stats: Optional[pd.DataFrame] = None
diff --git a/agentlib_mpc/utils/sampling.py b/agentlib_mpc/utils/sampling.py
index 90c37e4..5483cf6 100644
--- a/agentlib_mpc/utils/sampling.py
+++ b/agentlib_mpc/utils/sampling.py
@@ -86,6 +86,7 @@ def sample(
             f"does not match target ({target_grid_length})."
         )
     if isinstance(trajectory, pd.Series):
+        trajectory = trajectory.dropna()
         source_grid = np.array(trajectory.index)
         values = trajectory.values
     elif isinstance(trajectory, dict):
diff --git a/examples/admm/admm_example_coordinator.py b/examples/admm/admm_example_coordinator.py
index 2295f69..0af7101 100644
--- a/examples/admm/admm_example_coordinator.py
+++ b/examples/admm/admm_example_coordinator.py
@@ -28,6 +28,9 @@
 from agentlib.utils import MultiProcessingBroker
 from agentlib.utils.multi_agent_system import LocalMASAgency
 
+from agentlib_mpc.utils.plotting import admm_dashboard
+from agentlib_mpc.utils.plotting.admm_dashboard import show_admm_dashboard
+from agentlib_mpc.utils.plotting.interactive import show_dashboard
 
 agent_configs = [
     # use MS discretization method
@@ -64,7 +67,12 @@ def plot(results, start_pred=0):
 
 
 def run_example(
-    until=3000, with_plots=True, start_pred=0, log_level=logging.INFO, cleanup=True
+    until=3000,
+    with_plots=True,
+    start_pred=0,
+    log_level=logging.INFO,
+    cleanup=True,
+    show_dashboard: bool = False,
 ):
     # Set the log-level
     logging.basicConfig(level=log_level)
@@ -94,6 +102,15 @@ def run_example(
     mas.run(until=until)
     results = mas.get_results(cleanup=cleanup)
 
+    if show_dashboard:
+        show_admm_dashboard(
+            data={
+                "Cooler": results["Cooler"]["admm_module"],
+                "CooledRoom": results["CooledRoom"]["admm_module"],
+            },
+            residuals=results["Coordinator"]["admm_coordinator"],
+        )
+
     if with_plots:
         plot(results, start_pred=start_pred)
     return results
@@ -104,6 +121,7 @@ def run_example(
         with_plots=True,
         until=1800,
         start_pred=0,
+        show_dashboard=False,
         cleanup=True,
         log_level=logging.INFO,
     )
diff --git a/examples/one_room_mpc/ann/simple_mpc_nn.py b/examples/one_room_mpc/ann/simple_mpc_nn.py
index 66d4edc..d51944d 100644
--- a/examples/one_room_mpc/ann/simple_mpc_nn.py
+++ b/examples/one_room_mpc/ann/simple_mpc_nn.py
@@ -21,7 +21,7 @@ def agent_configs(ml_model_path: str) -> list[dict]:
                 "module_id": "myMPC",
                 "type": "agentlib_mpc.mpc",
                 "optimization_backend": {
-                    "type": "casadi_nn",
+                    "type": "casadi_ml",
                     "model": {
                         "type": {
                             "file": "model.py",
diff --git a/examples/one_room_mpc/gpr/simple_mpc_gpr.py b/examples/one_room_mpc/gpr/simple_mpc_gpr.py
index abe916b..b089f94 100644
--- a/examples/one_room_mpc/gpr/simple_mpc_gpr.py
+++ b/examples/one_room_mpc/gpr/simple_mpc_gpr.py
@@ -23,7 +23,7 @@ def agent_configs(ml_model_path: str) -> list[dict]:
                 "module_id": "myMPC",
                 "type": "agentlib_mpc.mpc",
                 "optimization_backend": {
-                    "type": "casadi_nn",
+                    "type": "casadi_ml",
                     "model": {
                         "type": {
                             "file": "model.py",
diff --git a/examples/one_room_mpc/linreg/simple_mpc_linreg.py b/examples/one_room_mpc/linreg/simple_mpc_linreg.py
index 76d7333..b543be0 100644
--- a/examples/one_room_mpc/linreg/simple_mpc_linreg.py
+++ b/examples/one_room_mpc/linreg/simple_mpc_linreg.py
@@ -23,7 +23,7 @@ def agent_configs(ml_model_path: str) -> list[dict]:
                 "module_id": "myMPC",
                 "type": "agentlib_mpc.mpc",
                 "optimization_backend": {
-                    "type": "casadi_nn",
+                    "type": "casadi_ml",
                     "model": {
                         "type": {
                             "file": "model.py",
diff --git a/examples/one_room_mpc/physical/simple_mpc.py b/examples/one_room_mpc/physical/simple_mpc.py
index ed2b105..c422e43 100644
--- a/examples/one_room_mpc/physical/simple_mpc.py
+++ b/examples/one_room_mpc/physical/simple_mpc.py
@@ -138,7 +138,7 @@ def setup_system(self):
                     "collocation_method": "legendre",
                 },
                 "solver": {
-                    "name": "ipopt",  # use fatrop with casadi 3.6.6 for speedup
+                    "name": "fatrop",  # use fatrop with casadi 3.6.6 for speedup
                 },
                 "results_file": "results//mpc.csv",
                 "save_results": True,
@@ -208,7 +208,7 @@ def run_example(
     )
     mas.run(until=until)
     try:
-        stats = load_mpc_stats("results/stats_mpc_fatropfull.csv")
+        stats = load_mpc_stats("results/__mpc.csv")
     except Exception:
         stats = None
     results = mas.get_results(cleanup=False)
@@ -267,5 +267,5 @@ def plot(mpc_results: pd.DataFrame, sim_res: pd.DataFrame, until: float):
 
 if __name__ == "__main__":
     run_example(
-        with_plots=True, with_dashboard=False, until=7200, log_level=logging.CRITICAL
+        with_plots=True, with_dashboard=True, until=7200, log_level=logging.INFO
     )
diff --git a/examples/output_ann/generate_training_data.py b/examples/output_ann/generate_training_data.py
index 212e2f0..0f3efe7 100644
--- a/examples/output_ann/generate_training_data.py
+++ b/examples/output_ann/generate_training_data.py
@@ -12,8 +12,8 @@
     CasadiOutput,
 )
 from agentlib_mpc.models.casadi_ml_model import CasadiMLModelConfig, CasadiMLModel
-from agentlib_mpc.models.serialized_ann import SerializedANN
-from agentlib_mpc.modules.ml_model_training.ann_trainer import ANNTrainer
+from agentlib_mpc.models.serialized_ml_model import SerializedANN
+from agentlib_mpc.modules.ml_model_training.ml_model_trainer import ANNTrainer
 
 
 def generate_test_data() -> pd.DataFrame:
diff --git a/pyproject.toml b/pyproject.toml
index de90040..bac641b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -43,7 +43,7 @@ dynamic = ["version"]
 
 fmu = ['FMPy>=0.3.6']
 ml = ["keras>=2.6.0", "tensorflow>=2.6.0", "scikit-learn"]
-interactive = ['plotly>=5.20.0', 'dash>=2.16.1']
+interactive = ['plotly>=5.20.0', 'dash>=2.16.1', 'dash_daq']
 full = ['agentlib_mpc[fmu,ml,interactive]', 'agentlib[full]']
 
 [package.urls]