MyFunctions_PIL.py

import numpy as np
import gymnasium as gym
from gymnasium import spaces
from typing import Dict, Any, Optional
import heyoka as hy
import time
import os
import torch

from stable_baselines3.common.monitor import Monitor
from stable_baselines3.common.utils import set_random_seed

def taylor_constant_thrust():
    """
    Build the Taylor-adaptive integrator for a constant-thrust spacecraft dynamics model.

    Author: Marcello Pareschi
    Last update: 2026-03-24

    This function defines the equations of motion of a spacecraft subject to
    central gravitational attraction and a constant low-thrust acceleration.
    The thrust magnitude and direction are treated as fixed parameters during
    propagation. The resulting Taylor-adaptive integrator can therefore be used
    to simulate trajectories under prescribed constant thrust conditions, which
    is useful for preliminary analyses, guidance studies, and validation
    purposes.

    Returns
    -------
    ta : heyoka.taylor_adaptive
        Taylor-adaptive integrator for the constant-thrust spacecraft dynamics.
    """

    # Define the state variables:
    # position (x, y, z), velocity (vx, vy, vz), and mass (m).
    x, y, z, vx, vy, vz, m = hy.make_vars("x", "y", "z", "vx", "vy", "vz", "m")

    # Compute ||r||^3 = (x^2 + y^2 + z^2)^(3/2),
    # required for the central gravitational acceleration term.
    r3 = (x**2 + y**2 + z**2) ** (1.5)

    # The thrust direction could alternatively be parameterized through
    # angular coordinates, as suggested in the commented formulation below.
    # ux = hy.par[2] * hy.cos(hy.par[3]) * hy.cos(hy.par[4])
    # uy = hy.par[2] * hy.sin(hy.par[3]) * hy.cos(hy.par[4])
    # uz = hy.par[2] * hy.sin(hy.par[4])

    # Define the dynamical system:
    # - position derivatives are equal to velocity,
    # - velocity derivatives include gravitational and thrust acceleration,
    # - mass derivative accounts for propellant consumption.
    syst = [
        (x, vx),
        (y, vy),
        (z, vz),
        (vx, -x / r3 + hy.par[0] * hy.par[2] * hy.par[3] / m),
        (vy, -y / r3 + hy.par[0] * hy.par[2] * hy.par[4] / m),
        (vz, -z / r3 + hy.par[0] * hy.par[2] * hy.par[5] / m),
        (m, -hy.par[0] / hy.par[1] * hy.par[2])
    ]

    # Set a strict numerical tolerance for high-accuracy propagation.
    tol_propagation = 1e-16

    # Define a default initial state for integrator construction.
    # These values can be overwritten before actual propagation.
    initial_state = [1.0] * 7

    # Define a default initial time.
    initial_time = 0.0

    # Define default parameter values:
    # these can be overwritten before actual propagation.
    initial_pars = [0.0] * 6

    # Build the Taylor-adaptive integrator for the constant-thrust system.
    ta = hy.taylor_adaptive(
        syst,
        state=initial_state,
        pars=initial_pars,
        time=initial_time,
        compact_mode=True,
        tol=tol_propagation
    )

    # Return the integrator.
    return ta

class LowThrustEnv(gym.Env):
    """
    Reinforcement-learning environment for low-thrust trajectory optimization.

    Author: Marcello Pareschi
    Last update: 2026-03-24

    This class implements a custom Gymnasium environment for the simulation
    and optimization of low-thrust spacecraft transfers. The environment
    propagates the spacecraft dynamics under thrust control, provides the
    agent with state-time observations, and evaluates the control policy
    through a reward function that combines fuel consumption, trajectory
    tracking, and terminal accuracy. Additional stochastic perturbations can
    be introduced on initial conditions, states, and controls in order to
    improve robustness during training.

    Parameters
    ----------
    config : dict
        Dictionary containing the full set of environment parameters, such as
        initial and final states, scaling constants, reward weights, time
        horizon, perturbation settings, and exploration options.
    """

    def __init__(self, config: Optional[Dict[str, Any]]) -> None:
        """
        Initialize the low-thrust environment.

        This constructor loads the configuration dictionary, defines the
        observation and action spaces, initializes the internal state, and
        creates the Taylor integrator used to propagate the spacecraft
        dynamics.
        """

        # Initialize the parent Gym environment.
        super(LowThrustEnv, self).__init__()

        # Load all configuration entries as object attributes.
        for key, item in config.items():
            setattr(self, key, item)

        # Define the time step size from the total time of flight
        # and the number of integration steps.
        self.dt = self.tof / self.N_steps

        # Define the observation space:
        # [position(3), velocity(3), mass(1), time(1)].
        low = np.array(
            [-self.rmax] * 3 + [-self.vmax] * 3 + [0.0] + [0.0],
            dtype=np.float32
        )
        high = np.array(
            [self.rmax] * 3 + [self.vmax] * 3 + [1.0] + [self.tof],
            dtype=np.float32
        )
        self.observation_space = spaces.Box(low, high, dtype=np.float32)

        # Define the action space:
        # [normalized throttle, normalized azimuth, normalized elevation].
        self.action_space = spaces.Box(
            np.array([-1] * 3),
            np.array([1] * 3),
            dtype=np.float32
        )

        # Initialize internal counters and the initial physical state.
        self.k = 0
        self.time = 0.0
        self.state = np.array([*self.r0, *self.v0, self.m0], dtype=np.float32)
        self.prev_state = self.state.copy()

        # Build the constant-thrust propagator and initialize its parameters.
        self.ta = taylor_constant_thrust()
        self.ta.pars[:2] = [self.c1, self.c2]
        self.ta.state[:] = self.state.copy()
        self.ta.time = 0.0

        # Counter used to track the total number of training steps.
        self.training_steps = 0

    def eps_constraint(self):
        """
        Return the current tolerance used for the terminal constraint.

        The tolerance is selected from a predefined schedule according to the
        current number of training steps. This allows the terminal constraint
        to be progressively tightened during training.

        Returns
        -------
        float
            Current admissible tolerance for the terminal constraint.
        """

        # Number of tolerance values available in the schedule.
        n_values = len(self.eps_schedule)

        # Number of training iterations assigned to each tolerance value.
        n_steps_per_value = self.N_iters / n_values

        # Select the appropriate tolerance value according to the current
        # training-step counter.
        for i in range(n_values):
            if self.training_steps <= (i + 1) * n_steps_per_value:
                return self.eps_schedule[i]

        # If all schedule thresholds have been exceeded, return the last value.
        return self.eps_schedule[n_values - 1]

    def get_observation(self, s):
        """
        Build the observation vector from the current state.

        The observation contains the physical state and the current time.
        In addition, the method updates the tracking error with respect to the
        reference trajectory at the current time step.

        Parameters
        ----------
        s : ndarray
            Current spacecraft state.

        Returns
        -------
        ndarray
            Observation vector composed of state and time.
        """

        # Create a local copy of the input state.
        state = s.copy()

        # Extract the reference state corresponding to the current step.
        state_ref = self.traj_ref[self.k, :]

        # Store the tracking error with respect to the reference trajectory.
        self.delta_ref = state - state_ref

        # Return the observation as [state, time].
        return np.array([*state, self.time], dtype=np.float32)

    def get_control(self, action):
        """
        Convert a normalized action vector into physical control variables.

        The first action component is mapped to a throttle value in [0, 1],
        while the remaining two components are interpreted as normalized
        azimuth and elevation angles and converted into a Cartesian thrust
        direction vector.

        Parameters
        ----------
        action : ndarray
            Normalized action vector.

        Returns
        -------
        ndarray
            Physical control vector [throttle, direction_x, direction_y, direction_z].
        """

        # Map the normalized throttle from [-1, 1] to [0, 1].
        throttle = 0.5 * (action[0] + 1)

        # Convert the normalized angular actions into physical angles.
        Az = action[1] * np.pi
        El = action[2] * (np.pi / 2)

        # Convert azimuth and elevation into Cartesian direction components.
        alpha = np.array([
            np.cos(El) * np.cos(Az),
            np.cos(El) * np.sin(Az),
            np.sin(El)
        ])

        # Return the full physical control vector.
        return np.array([throttle, *alpha], dtype=np.float32)

    def perturb_initial_conditions(self, s):
        """
        Perturb the initial conditions for exploration purposes.

        Gaussian noise is independently added to the position and velocity
        components of the state.

        Parameters
        ----------
        s : ndarray
            Initial state vector.

        Returns
        -------
        ndarray
            Perturbed initial state.
        """

        # Copy the state to avoid modifying the original input.
        s = s.copy()

        # Add Gaussian perturbations to the initial position.
        s[:3] += np.random.normal(0, self.sigma_ic_position, size=3)

        # Add Gaussian perturbations to the initial velocity.
        s[3:6] += np.random.normal(0, self.sigma_ic_velocity, size=3)

        # Return the perturbed initial state.
        return s

    def perturb_control(self, c):
        """
        Perturb the control input for exploration purposes.

        The throttle is perturbed by additive Gaussian noise and clipped to
        the admissible range [0, 1]. The thrust direction is perturbed by a
        small random rotation and re-normalized. If a mandatory thrust-off
        event is active, the throttle is set to zero within the specified
        time window.

        Parameters
        ----------
        c : ndarray
            Physical control vector.

        Returns
        -------
        ndarray
            Perturbed physical control vector.
        """

        # Create a local copy of the control input.
        control_unp = c.copy()
        throttle_unp, alpha_unp = control_unp[0], control_unp[1:]

        # Perturb the throttle and clip it to the admissible interval.
        throttle = np.clip(
            throttle_unp + np.random.normal(0, self.sigma_u_throttle),
            0.0,
            1.0
        )

        # Generate a small random rotation vector.
        rot = np.random.normal(0, self.sigma_u_direction, size=3)

        # Build a first-order rotation matrix approximation.
        A = np.array([
            [1, -rot[2], rot[1]],
            [rot[2], 1, -rot[0]],
            [-rot[1], rot[0], 1]
        ])

        # Rotate and re-normalize the thrust direction.
        alpha = A @ alpha_unp
        alpha /= np.linalg.norm(alpha)

        # If a mandatory thrust-off event is active, force zero throttle
        # in the corresponding time window.
        if self.mte_window is not None:
            start, end = self.mte_window
            if start <= self.k < end:
                throttle = 0.0

        # Return the perturbed control vector.
        return np.array([throttle, *alpha], dtype=np.float32)

    def perturb_state(self, s):
        """
        Perturb the propagated state for exploration purposes.

        Gaussian noise is independently added to the position and velocity
        components of the state.

        Parameters
        ----------
        s : ndarray
            Propagated state vector.

        Returns
        -------
        ndarray
            Perturbed state vector.
        """

        # Copy the state to avoid modifying the original input.
        s = s.copy()

        # Add Gaussian perturbations to the position.
        s[:3] += np.random.normal(0, self.sigma_s_position, size=3)

        # Add Gaussian perturbations to the velocity.
        s[3:6] += np.random.normal(0, self.sigma_s_velocity, size=3)

        # Return the perturbed state.
        return s

    def propagation_step(self, state, control, time):
        """
        Propagate the spacecraft dynamics for one time step.

        Parameters
        ----------
        state : ndarray
            Current spacecraft state.
        control : ndarray
            Physical control vector.
        time : float
            Current propagation time.

        Returns
        -------
        state : ndarray
            Propagated spacecraft state at the next time step.
        time : float
            Updated propagation time.
        """

        # Assign the control parameters to the propagator.
        self.ta.pars[2:] = control

        # Set the current propagation time.
        self.ta.time = time

        # Set the current state of the propagator.
        self.ta.state[:] = state

        # Propagate the dynamics until the next grid point.
        self.ta.propagate_until(self.time_grid[self.k + 1])

        # Increment the discrete step counter.
        self.k += 1

        # Store the propagated state.
        self.state = self.ta.state.copy()

        # Return the updated state and propagation time.
        return self.state, float(self.ta.time)

    def _terminated(self):
        """
        Check whether the episode has reached its natural terminal step.

        Returns
        -------
        bool
            True if the maximum number of steps has been reached.
        """

        return self.k >= self.N_steps

    def _truncated(self, state):
        """
        Check whether the episode must be truncated due to a constraint violation.

        Parameters
        ----------
        state : ndarray
            Current spacecraft state.

        Returns
        -------
        bool
            True if the state violates radius, velocity, or mass constraints.
        """

        return bool(
            np.linalg.norm(state[:3]) <= self.rmin or
            np.linalg.norm(state[:3]) >= self.rmax or
            np.linalg.norm(state[3:6]) >= self.vmax or
            state[-1] <= 0
        )

    def get_reward(self, observation, terminated, truncated):
        """
        Compute the reward associated with the current transition.

        The reward includes:
        - a fuel-consumption penalty,
        - a tracking penalty with respect to the reference trajectory,
        - a terminal penalty based on final-state violation or truncation.

        Parameters
        ----------
        observation : ndarray
            Current observation vector.
        terminated : bool
            Flag indicating whether the episode ended naturally.
        truncated : bool
            Flag indicating whether the episode was interrupted by a violation.

        Returns
        -------
        reward : float
            Scalar reward assigned to the transition.
        done : bool
            Boolean flag indicating whether the episode is over.
        """

        # Compute the fuel-consumption penalty from the mass variation.
        m_prev, m = self.prev_state[-1], self.state[-1]
        reward = -self.w_fuel * (m_prev - m)

        # Add the trajectory-tracking penalty.
        reward -= self.w_tracking * np.linalg.norm(self.delta_ref)

        # Determine whether the episode is over.
        done = terminated or truncated

        # Apply terminal penalties if the episode has ended.
        if done:
            # Relative terminal position violation.
            self.r_viol = np.linalg.norm(self.state[:3] - self.rf) / np.linalg.norm(self.rf)

            # Relative terminal velocity violation.
            self.v_viol = np.linalg.norm(self.state[3:6] - self.vf) / np.linalg.norm(self.vf)

            # Retrieve the current admissible tolerance.
            eps = self.eps_constraint()

            # Constraint violation beyond the current tolerance.
            c_viol = max(0.0, max(self.r_viol, self.v_viol) - eps)

            # Penalize truncation more severely, otherwise penalize terminal mismatch.
            if truncated:
                reward -= self.w_truncated
            else:
                reward -= self.w_terminated * c_viol

        # Return the final reward and the done flag.
        return float(reward), done

    def get_info(self, control, done):
        """
        Build the info dictionary used for logging and diagnostics.

        Parameters
        ----------
        control : ndarray
            Control applied at the current step.
        done : bool
            Flag indicating whether the episode has ended.

        Returns
        -------
        dict
            Dictionary containing step-wise and terminal diagnostic data.
        """

        # Build the step-wise logging dictionary.
        info = {
            "episode_step_data": {
                "r": [self.prev_state[:3]],
                "v": [self.prev_state[3:6]],
                "m": [self.prev_state[-1]],
                "t": [self.time],
                "u": [control],
            }
        }

        # If the episode is over, append terminal metrics and final state data.
        if done:
            info.update({
                "custom_metrics": {
                    "r_viol": self.r_viol,
                    "v_viol": self.v_viol,
                },
                "episode_end_data": {
                    "rf": self.state[:3],
                    "vf": self.state[3:6],
                    "mf": self.state[-1],
                }
            })

        # Return the full info dictionary.
        return info

    def reset(self, seed=None, options=None):
        """
        Reset the environment at the beginning of a new episode.

        The state is reinitialized, optional perturbations on the initial
        conditions are applied, and a possible mandatory thrust-off event is
        sampled if enabled.

        Parameters
        ----------
        seed : int, optional
            Random seed passed to the parent environment reset.
        options : dict, optional
            Additional reset options, currently unused.

        Returns
        -------
        obs : ndarray
            Initial observation of the new episode.
        info : dict
            Initial diagnostic dictionary.
        """

        # Reset the parent Gym environment.
        super().reset(seed=seed)

        # Reset the discrete time index and physical time.
        self.k = 0
        self.time = 0.0

        # Restore the nominal initial state.
        self.state = np.array([*self.r0, *self.v0, self.m0], dtype=np.float32)

        # Optionally perturb the initial conditions.
        if self.perturb_ics:
            self.state = self.perturb_initial_conditions(self.state)

        # By default, no mandatory thrust-off event is active.
        self.mte_window = None

        # Optionally sample a mandatory thrust-off window.
        if self.mte and np.random.rand() < self.pr_mte:
            duration = np.random.randint(self.n_mte_min, self.n_mte_max + 1)
            last_possible_start = self.N_steps - duration - 1

            if last_possible_start > 0:
                start_step = np.random.randint(0, last_possible_start)
                self.mte_window = (start_step, start_step + duration)

        # Store the previous state for reward computation.
        self.prev_state = self.state.copy()

        # Build the initial observation and the initial info dictionary.
        obs = self.get_observation(self.state)
        info = self.get_info(np.zeros(3), done=False)

        # Return the initial observation and info.
        return obs, info

    def step(self, action):
        """
        Advance the environment by one step using the provided action.

        The action is converted into a physical control, optionally perturbed,
        then used to propagate the spacecraft dynamics. The new observation,
        reward, termination flags, and diagnostic information are returned.

        Parameters
        ----------
        action : ndarray
            Action selected by the agent.

        Returns
        -------
        obs : ndarray
            New observation after propagation.
        reward : float
            Reward associated with the transition.
        terminated : bool
            True if the episode ended because the horizon was reached.
        truncated : bool
            True if the episode ended because a constraint was violated.
        info : dict
            Additional diagnostic information.
        """

        # Ensure that the action belongs to the admissible action space.
        assert self.action_space.contains(action)

        # Store the previous state before propagation.
        self.prev_state = self.state.copy()

        # Convert the normalized action into physical control variables.
        control = self.get_control(action)

        # Optionally perturb the control for exploration.
        if self.perturb_controls:
            control = self.perturb_control(control)

        # Propagate the system dynamics by one step.
        self.state, self.time = self.propagation_step(self.state, control, self.time)

        # Copy the propagated state and optionally perturb it before observation.
        s = self.state.copy()
        if self.perturb_states:
            s = self.perturb_state(s)

        # Build the new observation.
        obs = self.get_observation(s)

        # Check natural termination and truncation conditions.
        terminated = self._terminated()
        truncated = self._truncated(self.state)

        # Compute reward and global done flag.
        reward, done = self.get_reward(obs, terminated, truncated)

        # Build the diagnostic info dictionary.
        info = self.get_info(control, done)

        # Update the global training-step counter.
        self.training_steps += 1

        # Return the full Gymnasium step tuple.
        return obs, reward, terminated, truncated, info
    
def run_pil_simulation(env, url, session):
    """
    Run a single Processor-in-the-Loop (PIL) simulation episode.

    Author: Marcello Pareschi
    Last update: 2026-03-24

    This function executes one complete PIL simulation episode by coupling a
    local dynamical environment with an external inference server. At each
    step, the current observation is sent to the remote device, which returns
    the control action together with inference and hardware-monitoring data.
    The environment is then advanced locally, and all relevant metrics are
    collected throughout the episode. The function finally returns a structured
    dictionary containing timing statistics, hardware usage, trajectory
    history, control history, and terminal mission-performance indicators.

    Parameters
    ----------
    env : gym.Env
        Environment used to simulate the spacecraft dynamics locally.
    url : str
        URL of the external inference endpoint.
    session : requests.Session
        Active HTTP session used to communicate with the external device.

    Returns
    -------
    dict
        Dictionary containing the collected episode data, including:
        - inference times,
        - network overheads,
        - local physics-integration times,
        - power consumption,
        - GPU utilization,
        - trajectory history,
        - control history,
        - terminal violations,
        - final state.
    """

    # Reset the environment and obtain the initial observation.
    obs = env.reset()

    # Initialize containers for timing and hardware metrics.
    inf_times = []          # Remote inference times
    net_overheads = []      # Communication overhead excluding inference
    physics_times = []      # Local environment propagation times
    pwr_cons = []           # Measured power consumption
    gpu_utils = []          # Measured GPU utilization

    # Initialize containers for the trajectory history.
    t_history = []          # Time history
    r_history = []          # Position history
    v_history = []          # Velocity history
    m_history = []          # Mass history
    controls = []           # Applied control history

    # Initialize terminal-violation variables.
    r_viol = None
    v_viol = None

    # Initialize placeholders for the final state.
    rf = None
    vf = None
    mf = None

    # Episode loop.
    done = False
    while not done:
        # --- PHASE 1: REMOTE INFERENCE ---
        # Send the current observation to the external inference server
        # and receive the predicted control action.
        comm_start = time.perf_counter()
        try:
            # Convert the observation to raw bytes for efficient transmission.
            obs_bytes = obs[0].astype(np.float32).tobytes()

            # Send the observation to the external server.
            response = session.post(
                url,
                data=obs_bytes,
                headers={"Content-Type": "application/octet-stream"},
                timeout=2
            )
            response.raise_for_status()

            # Decode the server response.
            res_data = response.json()

            # Extract the control action and inference-related metrics.
            action = res_data["action"]
            jetson_inf = res_data["inference_time_ms"]
            gpu_u = res_data.get("gpu_utilization_percent", 0)
            pwr = res_data.get("power_core_mw", 0)

        except Exception as e:
            # Stop the episode if communication with the inference server fails.
            print(f"Communication Error during episode: {e}")
            break

        # Compute the total round-trip communication time.
        total_comm_ms = (time.perf_counter() - comm_start) * 1000

        # Estimate the network overhead by subtracting the inference time
        # reported by the remote device from the total communication time.
        net_overhead = max(0, total_comm_ms - jetson_inf)

        # --- PHASE 2: LOCAL PHYSICS PROPAGATION ---
        # Advance the local environment using the action returned
        # by the external inference server.
        physics_start = time.perf_counter()
        obs, reward, dones, infos = env.step([action])
        physics_ms = (time.perf_counter() - physics_start) * 1000

        # --- DATA COLLECTION ---
        # Store timing and hardware-monitoring metrics for the current step.
        inf_times.append(jetson_inf)
        net_overheads.append(net_overhead)
        physics_times.append(physics_ms)
        pwr_cons.append(pwr)
        gpu_utils.append(gpu_u)

        # Extract the step-wise diagnostic data from the environment info.
        info = infos[0]
        step_data = info["episode_step_data"]

        # Append the current step data to the history containers.
        # Each quantity is extracted from a single-element list.
        t_history.append(step_data["t"][0])
        r_history.append(step_data["r"][0])
        v_history.append(step_data["v"][0])
        m_history.append(step_data["m"][0])
        controls.append(step_data["u"][0])

        # Update the termination flag.
        done = dones[0]

        # --- PHASE 3: TERMINAL METRICS ---
        # If the episode has ended, extract the final state
        # and the terminal constraint violations.
        if done:
            # Retrieve the final state reached at the end of the episode.
            end_data = info["episode_end_data"]

            # Append the terminal state to the trajectory history.
            r_history.append(end_data["rf"])
            v_history.append(end_data["vf"])
            m_history.append(end_data["mf"])

            # Extract the terminal position and velocity violations.
            r_viol = info["custom_metrics"]["r_viol"]
            v_viol = info["custom_metrics"]["v_viol"]

            # Store the final state explicitly.
            rf = end_data["rf"]
            vf = end_data["vf"]
            mf = end_data["mf"]

    # Assemble all collected episode data into a structured dictionary.
    ep_data = {
        "inf_times": inf_times,
        "net_overheads": net_overheads,
        "physics_times": physics_times,
        "pwr_cons": pwr_cons,
        "gpu_utils": gpu_utils,
        "steps": len(t_history),
        "t_history": t_history,
        "r_history": r_history,
        "v_history": v_history,
        "m_history": m_history,
        "controls": controls,
        "r_viol": r_viol,
        "v_viol": v_viol,
        "rf": rf,
        "vf": vf,
        "mf": mf
    }

    # Return the complete set of episode metrics.
    return ep_data

def make_env(config: Dict[str, Any], rank: int, seed: int = 0) -> gym.Env:
    """
    Create a callable environment factory for parallel low-thrust RL simulations.

    Author: Marcello Pareschi
    Last update: 2026-03-24

    This function returns an initializer that builds a single instance of the
    low-thrust environment and applies a worker-specific random seed. The
    returned callable is especially useful when constructing vectorized
    environments for parallel reinforcement-learning training or evaluation.

    Parameters
    ----------
    config : dict
        Dictionary containing the configuration parameters required to
        initialize the low-thrust environment.
    rank : int
        Identifier of the worker environment. It is used to generate a
        distinct random seed for each parallel instance.
    seed : int, optional
        Base random seed shared across workers.

    Returns
    -------
    callable
        Environment-construction function to be used by vectorized or
        parallel reinforcement-learning wrappers.
    """

    def _init() -> gym.Env:
        """
        Build and return one environment instance.

        Returns
        -------
        gym.Env
            Initialized low-thrust environment.
        """

        # Create a new low-thrust environment using the provided configuration.
        env = LowThrustEnv(config)

        # Reset the environment with a worker-specific random seed.
        env.reset(seed=seed + rank)

        # Return the initialized environment instance.
        return env

    # Set the global random seed for reproducibility.
    set_random_seed(seed)

    # Return the environment factory function.
    return _init

def run_pil_simulation_batch(env, url, session):
    """
    Run a batch of Processor-in-the-Loop (PIL) simulation episodes in parallel.

    Author: Marcello Pareschi
    Last update: 2026-03-24

    This function executes multiple PIL simulation episodes simultaneously by
    coupling a vectorized local environment with an external inference server.
    At each iteration, the full batch of observations is sent to the remote
    device, which returns a batch of control actions together with aggregate
    timing and hardware metrics. The local environments are then propagated in
    parallel, and the trajectory history, timing data, and terminal metrics are
    collected separately for each environment in the batch. The function
    returns one structured dictionary per completed episode.

    Parameters
    ----------
    env : gym.Env
        Vectorized environment containing multiple low-thrust simulation
        instances.
    url : str
        URL of the external inference endpoint.
    session : requests.Session
        Active HTTP session used to communicate with the external device.

    Returns
    -------
    list of dict
        List containing one dictionary for each environment in the batch. Each
        dictionary stores timing metrics, trajectory history, control history,
        and terminal performance indicators.
    """

    # Retrieve the number of parallel environments in the batch.
    n_envs = env.num_envs

    # Reset the vectorized environment and obtain the initial observations.
    obs = env.reset()

    # Initialize one buffer per metric and per environment.
    # Each entry is a list that will store the corresponding time history
    # for a single trajectory in the batch.
    batch_metrics = {
        "inf_times": [[] for _ in range(n_envs)],
        "net_overheads": [[] for _ in range(n_envs)],
        "physics_times": [[] for _ in range(n_envs)],
        "pwr_cons": [[] for _ in range(n_envs)],
        "t_hist": [[] for _ in range(n_envs)],
        "r_hist": [[] for _ in range(n_envs)],
        "v_hist": [[] for _ in range(n_envs)],
        "m_hist": [[] for _ in range(n_envs)],
        "controls": [[] for _ in range(n_envs)]
    }

    # Track which environments have already completed their episode.
    dones = np.zeros(n_envs, dtype=bool)

    # Placeholder for the final structured output of each environment.
    episode_results = [None] * n_envs

    # Continue until all environments in the batch have terminated.
    while not np.all(dones):
        # --- PHASE 1: REMOTE BATCH INFERENCE ---
        # Send the full batch of observations to the external inference server.
        comm_start = time.perf_counter()
        try:
            # Convert the full observation batch into raw bytes.
            obs_bytes = obs.astype(np.float32).tobytes()

            # Send the observation batch to the remote inference endpoint.
            response = session.post(
                url,
                data=obs_bytes,
                headers={"Content-Type": "application/octet-stream"},
                timeout=5
            )
            response.raise_for_status()

            # Decode the response returned by the remote device.
            res_data = response.json()

            # Extract the predicted actions for all environments.
            actions = res_data["action"]

            # Extract the average inference time for the full batch.
            jetson_inf = res_data["inference_time_ms"]

            # Extract the reported power consumption.
            pwr = res_data.get("power_core_mw", 0)

        except Exception as e:
            # Abort the batch simulation if communication fails.
            print(f"Communication Error during batch: {e}")
            break

        # Compute the total round-trip communication time.
        total_comm_ms = (time.perf_counter() - comm_start) * 1000

        # Estimate the communication overhead excluding remote inference time.
        net_overhead = max(0, total_comm_ms - jetson_inf)

        # --- PHASE 2: LOCAL BATCH PROPAGATION ---
        # Advance all local environments in parallel using the returned actions.
        physics_start = time.perf_counter()
        obs, rewards, step_dones, infos = env.step(actions)
        physics_ms = (time.perf_counter() - physics_start) * 1000

        # --- DATA COLLECTION FOR EACH ENVIRONMENT ---
        for i in range(n_envs):
            # Collect data only for environments that were still active
            # at the start of this iteration.
            if not dones[i]:
                info = infos[i]
                step_data = info["episode_step_data"]

                # Store timing and power-related metrics.
                batch_metrics["inf_times"][i].append(jetson_inf)
                batch_metrics["net_overheads"][i].append(net_overhead)
                batch_metrics["physics_times"][i].append(physics_ms)
                batch_metrics["pwr_cons"][i].append(pwr)

                # Store the trajectory and control histories.
                batch_metrics["t_hist"][i].append(step_data["t"][0])
                batch_metrics["r_hist"][i].append(step_data["r"][0])
                batch_metrics["v_hist"][i].append(step_data["v"][0])
                batch_metrics["m_hist"][i].append(step_data["m"][0])
                batch_metrics["controls"][i].append(step_data["u"][0])

                # If this specific environment has just terminated,
                # finalize and store its episode results.
                if step_dones[i]:
                    end_data = info["episode_end_data"]

                    # Append the final state reached at the end of the episode.
                    batch_metrics["r_hist"][i].append(end_data["rf"])
                    batch_metrics["v_hist"][i].append(end_data["vf"])
                    batch_metrics["m_hist"][i].append(end_data["mf"])

                    # Build the final structured output for this environment.
                    episode_results[i] = {
                        "inf_times": batch_metrics["inf_times"][i],
                        "net_overheads": batch_metrics["net_overheads"][i],
                        "physics_times": batch_metrics["physics_times"][i],
                        "pwr_cons": batch_metrics["pwr_cons"][i],
                        "t_history": batch_metrics["t_hist"][i],
                        "r_history": batch_metrics["r_hist"][i],
                        "v_history": batch_metrics["v_hist"][i],
                        "m_history": batch_metrics["m_hist"][i],
                        "controls": batch_metrics["controls"][i],
                        "r_viol": info["custom_metrics"]["r_viol"],
                        "v_viol": info["custom_metrics"]["v_viol"],
                        "rf": end_data["rf"],
                        "vf": end_data["vf"],
                        "mf": end_data["mf"],
                        "steps": len(batch_metrics["t_hist"][i])
                    }

        # Update the done flags by marking as completed all environments
        # that terminated at the current iteration.
        dones = np.logical_or(dones, step_dones)

    # Return the list of episode results, one for each environment.
    return episode_results

def run_sil_simulation(env, ort_sess):
    """
    Run a single Software-in-the-Loop (SIL) simulation episode.

    Author: Marcello Pareschi
    Last update: 2026-03-24

    This function executes one complete SIL simulation episode by coupling a
    local dynamical environment with an ONNX Runtime inference session. At each
    step, the current observation is processed by the neural network, the
    predicted action is clipped to the admissible range, and the environment is
    advanced locally. During the episode, inference times, local physics times,
    trajectory history, control history, and terminal performance indicators
    are collected. The function returns all relevant episode data in a
    structured dictionary.

    Parameters
    ----------
    env : gym.Env
        Environment used to simulate the spacecraft dynamics locally.
    ort_sess : onnxruntime.InferenceSession
        ONNX Runtime session used to perform policy inference.

    Returns
    -------
    dict
        Dictionary containing the collected episode data, including:
        - inference times,
        - local physics-integration times,
        - trajectory history,
        - control history,
        - terminal violations,
        - final state.
    """

    # Reset the environment and obtain the initial observation.
    obs = env.reset()

    # Initialize containers for timing metrics.
    inf_times = []          # ONNX Runtime inference times
    physics_times = []      # Local environment propagation times

    # Initialize containers for the trajectory history.
    t_history = []          # Time history
    r_history = []          # Position history
    v_history = []          # Velocity history
    m_history = []          # Mass history
    controls = []           # Applied control history

    # Initialize placeholders for terminal metrics.
    r_viol = None
    v_viol = None
    rf = None
    vf = None
    mf = None

    # Episode loop.
    done = False
    while not done:
        # --- PHASE 1: LOCAL ONNX INFERENCE ---
        # Run policy inference with ONNX Runtime on the current observation.
        start_inf = time.perf_counter()
        outputs = ort_sess.run(None, {"observation": obs})[0]
        inference_time = (time.perf_counter() - start_inf) * 1000

        # Extract the action from the network output,
        # flatten it, and clip it to the admissible action range [-1, 1].
        action = outputs[0].flatten()
        action_clipped = np.clip(action, -1.0, 1.0).tolist()

        # --- PHASE 2: LOCAL PHYSICS PROPAGATION ---
        # Advance the environment using the clipped action.
        physics_start = time.perf_counter()
        obs, reward, dones, infos = env.step([action_clipped])
        physics_ms = (time.perf_counter() - physics_start) * 1000

        # Extract the step-wise diagnostic data from the environment info.
        info = infos[0]
        step_data = info["episode_step_data"]

        # Store timing metrics for the current step.
        inf_times.append(inference_time)
        physics_times.append(physics_ms)

        # Store the trajectory and control histories.
        t_history.append(step_data["t"][0])
        r_history.append(step_data["r"][0])
        v_history.append(step_data["v"][0])
        m_history.append(step_data["m"][0])
        controls.append(step_data["u"][0])

        # Update the termination flag.
        done = dones[0]

        # --- PHASE 3: TERMINAL METRICS ---
        # If the episode has ended, extract the final state
        # and the terminal constraint violations.
        if done:
            end_data = info["episode_end_data"]

            # Append the final state reached at the end of the episode.
            r_history.append(end_data["rf"])
            v_history.append(end_data["vf"])
            m_history.append(end_data["mf"])

            # Extract terminal position and velocity violations.
            r_viol = info["custom_metrics"]["r_viol"]
            v_viol = info["custom_metrics"]["v_viol"]

            # Store the final state explicitly.
            rf = end_data["rf"]
            vf = end_data["vf"]
            mf = end_data["mf"]

    # Assemble all collected episode data into a structured dictionary.
    ep_data = {
        "inf_times": inf_times,
        "physics_times": physics_times,
        "steps": len(t_history),
        "t_history": t_history,
        "r_history": r_history,
        "v_history": v_history,
        "m_history": m_history,
        "controls": controls,
        "r_viol": r_viol,
        "v_viol": v_viol,
        "rf": rf,
        "vf": vf,
        "mf": mf
    }

    # Return the complete set of episode metrics.
    return ep_data

def _to_torch_obs(obs, device, dtype):
    """
    Convert an environment observation into a PyTorch tensor with batch dimension.

    Author: Marcello Pareschi
    Last update: 2026-03-24

    This utility function transforms an observation coming from a Gym or
    Gymnasium environment into a PyTorch tensor compatible with neural-network
    inference. If the observation is one-dimensional, a batch dimension is
    added so that the output tensor has shape (1, obs_dim). The tensor is then
    moved to the specified device and converted to the desired floating-point
    precision.

    Parameters
    ----------
    obs : array-like
        Environment observation to be converted.
    device : torch.device
        Target device on which the tensor will be allocated.
    dtype : torch.dtype
        Target tensor data type.

    Returns
    -------
    torch.Tensor
        Observation converted to a PyTorch tensor with an explicit batch
        dimension.
    """

    # Convert the observation to a NumPy array if it is not already one.
    obs_np = np.asarray(obs)

    # If the observation is one-dimensional, add a batch dimension so that
    # it becomes compatible with batched neural-network inference.
    if obs_np.ndim == 1:
        obs_np = obs_np[None, :]  # Shape becomes (1, obs_dim)

    # Convert the NumPy array to a PyTorch tensor and move it
    # to the requested device and dtype.
    return torch.from_numpy(obs_np).to(device=device, dtype=dtype)

def _torch_infer(model, obs_t, device):
    """
    Perform a forward pass with a PyTorch model and measure inference time.

    Author: Marcello Pareschi
    Last update: 2026-03-24

    This utility function executes a neural-network forward pass on the
    provided observation tensor while measuring the corresponding inference
    latency. If the selected device is CUDA-enabled, synchronization is
    performed before and after the forward pass in order to obtain an accurate
    timing measurement. The output is converted back to a NumPy array and
    reshaped into a one-dimensional action vector.

    Parameters
    ----------
    model : nn.Module
        PyTorch model used for inference.
    obs_t : torch.Tensor
        Input observation tensor.
    device : torch.device
        Device on which the inference is executed.

    Returns
    -------
    action_np : ndarray
        Predicted action converted to a one-dimensional NumPy array.
    inference_ms : float
        Measured inference time in milliseconds.
    """

    # Synchronize the CUDA device before timing in order to avoid
    # including queued operations from previous kernels.
    if device.type == "cuda":
        torch.cuda.synchronize()

    # Start the high-resolution timer.
    t0 = time.perf_counter()

    # Run the forward pass without gradient tracking,
    # since this function is intended for inference only.
    with torch.no_grad():
        out = model(obs_t)

    # Synchronize again after inference to ensure that all GPU operations
    # have completed before stopping the timer.
    if device.type == "cuda":
        torch.cuda.synchronize()

    # Compute the elapsed inference time in milliseconds.
    inference_ms = (time.perf_counter() - t0) * 1000.0

    # Move the output to CPU and convert it to a NumPy array.
    out_np = out.detach().cpu().numpy()

    # If the output is batched, extract the first element;
    # otherwise keep the array as it is.
    action_np = out_np[0] if out_np.ndim == 2 else out_np

    # Return the action as a flattened NumPy vector together with the timing.
    return np.asarray(action_np).reshape(-1), inference_ms

def _reset_env(env):
    """
    Reset an environment with Gym/Gymnasium compatibility handling.

    Author: Marcello Pareschi
    Last update: 2026-03-24

    This utility function resets an environment while handling the difference
    between the Gym and Gymnasium reset interfaces. In Gymnasium, the reset
    method typically returns a tuple `(obs, info)`, whereas in older Gym
    versions it may return only the observation. This function extracts and
    returns the observation in a unified way.

    Parameters
    ----------
    env : gym.Env
        Environment to be reset.

    Returns
    -------
    obs : array-like
        Initial observation returned by the environment reset.
    """

    # Reset the environment.
    out = env.reset()

    # Handle the Gymnasium interface, where reset returns (obs, info).
    if isinstance(out, tuple) and len(out) == 2:
        obs, _info = out
    else:
        # Handle the older Gym interface, where reset returns only the observation.
        obs = out

    # Return the initial observation.
    return obs

def run_sil_simulation_torch(
    env,
    model,
    precision="float64",
    device=None,
    clip_low=-1.0,
    clip_high=1.0,
    warmup_steps=50,
):
    """
    Run a single Software-in-the-Loop (SIL) simulation episode using a PyTorch model.

    Author: Marcello Pareschi
    Last update: 2026-03-24

    This function performs a closed-loop SIL simulation in which the control
    action is generated online by a PyTorch neural network. At each step, the
    current environment observation is converted into a tensor, processed by
    the model, clipped to the admissible action range, and then applied to the
    environment. The function records inference times, local physics-step
    times, trajectory history, control history, and terminal mission metrics.
    It also supports configurable numerical precision and optional warm-up
    iterations before the actual simulation begins.

    Parameters
    ----------
    env : gym.Env
        Environment used to simulate the spacecraft dynamics locally.
    model : nn.Module
        PyTorch model used for policy inference.
    precision : str, optional
        Numerical precision used for inference. Supported values are
        'float64', 'float32', and 'float16'.
    device : torch.device, optional
        Device used for inference. If not provided, CUDA is selected when
        available, otherwise CPU is used.
    clip_low : float, optional
        Lower bound used to clip the predicted action.
    clip_high : float, optional
        Upper bound used to clip the predicted action.
    warmup_steps : int, optional
        Number of warm-up forward passes performed before the actual episode,
        useful for reducing first-inference overhead.

    Returns
    -------
    dict
        Dictionary containing the collected episode data, including:
        - inference times,
        - local physics-step times,
        - trajectory history,
        - control history,
        - terminal violations,
        - final state,
        - inference precision and device information.
    """

    # If no device is explicitly provided, automatically select CUDA when
    # available; otherwise fall back to CPU.
    if device is None:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # Set the numerical precision of the model according to the user request.
    if precision == "float64":
        dtype = torch.float64
        model = model.double()
    elif precision == "float32":
        dtype = torch.float32
        model = model.float()
    elif precision == "float16":
        dtype = torch.float16
        model = model.half()
    else:
        raise ValueError(
            "Unsupported precision. Use 'float64', 'float32', or 'float16'."
        )

    # Move the model to the selected device and switch it to evaluation mode.
    model = model.to(device)
    model.eval()

    # Reset the environment and obtain the initial observation.
    obs = _reset_env(env)

    # Optional warm-up phase:
    # perform a number of forward passes without stepping the environment in
    # order to reduce cold-start effects during the actual timing collection.
    for _ in range(max(0, warmup_steps)):
        obs_t = _to_torch_obs(obs, device, dtype)
        _ = model(obs_t)

    # Initialize containers for timing metrics.
    inf_times = []          # PyTorch inference times
    physics_times = []      # Local environment propagation times

    # Initialize containers for the trajectory history.
    t_history = []          # Time history
    r_history = []          # Position history
    v_history = []          # Velocity history
    m_history = []          # Mass history
    controls = []           # Applied control history

    # Initialize placeholders for terminal metrics.
    done = False
    r_viol = v_viol = None
    rf = vf = mf = None

    # Episode loop.
    while not done:
        # --- PHASE 1: PYTORCH INFERENCE ---
        # Convert the current observation into a PyTorch tensor and run the
        # model inference while measuring the elapsed time.
        obs_t = _to_torch_obs(obs, device, dtype)
        action, inference_ms = _torch_infer(model, obs_t, device)

        # Clip the predicted action to the admissible action range and convert
        # it to a standard Python list for environment compatibility.
        action_clipped = np.clip(action, clip_low, clip_high).tolist()

        # --- PHASE 2: LOCAL PHYSICS PROPAGATION ---
        # Advance the environment using the clipped control action.
        t0 = time.perf_counter()
        obs, reward, dones, infos = env.step([action_clipped])
        physics_ms = (time.perf_counter() - t0) * 1000.0

        # Extract the diagnostic information returned by the environment.
        info = infos[0]
        step_data = info["episode_step_data"]

        # Store timing metrics for the current step.
        inf_times.append(inference_ms)
        physics_times.append(physics_ms)

        # Store the trajectory and control histories.
        t_history.append(step_data["t"][0])
        r_history.append(step_data["r"][0])
        v_history.append(step_data["v"][0])
        m_history.append(step_data["m"][0])
        controls.append(step_data["u"][0])

        # Update the termination flag.
        done = bool(dones[0])

        # --- PHASE 3: TERMINAL METRICS ---
        # If the episode has ended, extract the final state and the terminal
        # constraint violations.
        if done:
            end_data = info["episode_end_data"]

            # Append the final state reached at the end of the episode.
            r_history.append(end_data["rf"])
            v_history.append(end_data["vf"])
            m_history.append(end_data["mf"])

            # Extract terminal position and velocity violations.
            r_viol = info["custom_metrics"]["r_viol"]
            v_viol = info["custom_metrics"]["v_viol"]

            # Store the final state explicitly.
            rf, vf, mf = end_data["rf"], end_data["vf"], end_data["mf"]

    # Assemble all collected episode data into a structured dictionary.
    ep_data = {
        "inf_times": inf_times,
        "physics_times": physics_times,
        "steps": len(t_history),
        "t_history": t_history,
        "r_history": r_history,
        "v_history": v_history,
        "m_history": m_history,
        "controls": controls,
        "r_viol": r_viol,
        "v_viol": v_viol,
        "rf": rf,
        "vf": vf,
        "mf": mf,
        "precision": precision,
        "dtype": str(dtype),
        "device": str(device),
    }

    # Return the complete set of episode metrics.
    return ep_data

def run_sil_simulation_torch1(
    env,
    model,
    precision="float64",
    device=None,
    clip_low=-1.0,
    clip_high=1.0,
    warmup_steps=50,
):
    """
    Run a single Software-in-the-Loop (SIL) simulation episode using a PyTorch model
    and store both observations and predicted actions.

    Author: Marcello Pareschi
    Last update: 2026-03-24

    This function performs a closed-loop SIL simulation in which the control
    action is generated online by a PyTorch neural network. At each step, the
    current environment observation is converted into a tensor, processed by
    the model, clipped to the admissible action range, and then applied to the
    environment. In addition to the standard trajectory, control, and timing
    metrics, this version also stores the full sequence of observations and raw
    predicted actions, which is useful for post-processing, debugging, and
    model-analysis tasks.

    Parameters
    ----------
    env : gym.Env
        Environment used to simulate the spacecraft dynamics locally.
    model : nn.Module
        PyTorch model used for policy inference.
    precision : str, optional
        Numerical precision used for inference. Supported values are
        'float64', 'float32', and 'float16'.
    device : torch.device, optional
        Device used for inference. If not provided, CUDA is selected when
        available, otherwise CPU is used.
    clip_low : float, optional
        Lower bound used to clip the predicted action.
    clip_high : float, optional
        Upper bound used to clip the predicted action.
    warmup_steps : int, optional
        Number of warm-up forward passes performed before the actual episode,
        useful for reducing first-inference overhead.

    Returns
    -------
    dict
        Dictionary containing the collected episode data, including:
        - inference times,
        - local physics-step times,
        - trajectory history,
        - control history,
        - terminal violations,
        - final state,
        - stored observations,
        - stored raw actions,
        - inference precision and device information.
    """

    # If no device is explicitly provided, automatically select CUDA when
    # available; otherwise fall back to CPU.
    if device is None:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # Set the numerical precision of the model according to the user request.
    if precision == "float64":
        dtype = torch.float64
        model = model.double()
    elif precision == "float32":
        dtype = torch.float32
        model = model.float()
    elif precision == "float16":
        dtype = torch.float16
        model = model.half()
    else:
        raise ValueError(
            "Unsupported precision. Use 'float64', 'float32', or 'float16'."
        )

    # Move the model to the selected device and switch it to evaluation mode.
    model = model.to(device)
    model.eval()

    # Reset the environment and obtain the initial observation.
    obs = _reset_env(env)

    # Optional warm-up phase:
    # perform a number of forward passes without stepping the environment in
    # order to reduce cold-start effects during the actual timing collection.
    for _ in range(max(0, warmup_steps)):
        obs_t = _to_torch_obs(obs, device, dtype)
        _ = model(obs_t)

    # Initialize containers for timing metrics.
    inf_times = []          # PyTorch inference times
    physics_times = []      # Local environment propagation times

    # Initialize containers for the trajectory history.
    t_history = []          # Time history
    r_history = []          # Position history
    v_history = []          # Velocity history
    m_history = []          # Mass history
    controls = []           # Applied control history

    # Initialize placeholders for terminal metrics.
    done = False
    r_viol = v_viol = None
    rf = vf = mf = None

    # Initialize containers for the raw observation and action histories.
    obs_list = []
    act_list = []

    # Episode loop.
    while not done:
        # --- PHASE 1: PYTORCH INFERENCE ---
        # Convert the current observation into a PyTorch tensor and run the
        # model inference while measuring the elapsed time.
        obs_t = _to_torch_obs(obs, device, dtype)
        action, inference_ms = _torch_infer(model, obs_t, device)

        # Clip the predicted action to the admissible action range and convert
        # it to a standard Python list for environment compatibility.
        action_clipped = np.clip(action, clip_low, clip_high).tolist()

        # Store the current observation and the raw predicted action.
        obs_list.append(obs.squeeze())
        act_list.append(action.squeeze())

        # --- PHASE 2: LOCAL PHYSICS PROPAGATION ---
        # Advance the environment using the clipped control action.
        t0 = time.perf_counter()
        obs, reward, dones, infos = env.step([action_clipped])
        physics_ms = (time.perf_counter() - t0) * 1000.0

        # Extract the diagnostic information returned by the environment.
        info = infos[0]
        step_data = info["episode_step_data"]

        # Store timing metrics for the current step.
        inf_times.append(inference_ms)
        physics_times.append(physics_ms)

        # Store the trajectory and control histories.
        t_history.append(step_data["t"][0])
        r_history.append(step_data["r"][0])
        v_history.append(step_data["v"][0])
        m_history.append(step_data["m"][0])
        controls.append(step_data["u"][0])

        # Update the termination flag.
        done = bool(dones[0])

        # --- PHASE 3: TERMINAL METRICS ---
        # If the episode has ended, extract the final state and the terminal
        # constraint violations.
        if done:
            end_data = info["episode_end_data"]

            # Append the final state reached at the end of the episode.
            r_history.append(end_data["rf"])
            v_history.append(end_data["vf"])
            m_history.append(end_data["mf"])

            # Extract terminal position and velocity violations.
            r_viol = info["custom_metrics"]["r_viol"]
            v_viol = info["custom_metrics"]["v_viol"]

            # Store the final state explicitly.
            rf, vf, mf = end_data["rf"], end_data["vf"], end_data["mf"]

    # Assemble all collected episode data into a structured dictionary.
    ep_data = {
        "inf_times": inf_times,
        "physics_times": physics_times,
        "steps": len(t_history),
        "t_history": t_history,
        "r_history": r_history,
        "v_history": v_history,
        "m_history": m_history,
        "controls": controls,
        "r_viol": r_viol,
        "v_viol": v_viol,
        "rf": rf,
        "vf": vf,
        "mf": mf,
        "precision": precision,
        "dtype": str(dtype),
        "device": str(device),
        "obs_list": obs_list,
        "act_list": act_list,
    }

    # Return the complete set of episode metrics.
    return ep_data