raddo 0.1.3__tar.gz

raddo-0.1.3/.gitignore

@@ -0,0 +1,9 @@
+__pycache__/
+__marimo__/
+c_generated_code/
+wandb/
+figures/
+*.json
+*.png
+.env
+.marimo.toml

raddo-0.1.3/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Patrick Schwartz
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

raddo-0.1.3/PKG-INFO

@@ -0,0 +1,142 @@
+Metadata-Version: 2.5
+Name: raddo
+Version: 0.1.3
+Summary: RApidly Adapting Data-Driven Opimization (RADDO). JAX-based trajectory generator with Moreau-based QP solver with an emphasis on adaptation and attitude dynamics
+Project-URL: Homepage, https://github.com/JHU-ACEL/raddo
+Author-email: Patrick Schwartz <pschwa24@jh.edu>
+Maintainer-email: Patrick Schwartz <pschwa24@jh.edu>
+License-Expression: MIT
+License-File: LICENSE
+Classifier: Programming Language :: Python
+Requires-Python: >=3.12
+Requires-Dist: equinox>=0.13.8
+Requires-Dist: jax>=0.11.2
+Requires-Dist: moreau>=0.4.1
+Requires-Dist: numpy>=2.5.3
+Requires-Dist: optax>=0.2.8
+Provides-Extra: cuda12
+Requires-Dist: jax[cuda12]>=0.11.2; extra == 'cuda12'
+Requires-Dist: moreau[cuda12]>=0.4.1; extra == 'cuda12'
+Provides-Extra: cuda13
+Requires-Dist: jax[cuda13]>=0.11.2; extra == 'cuda13'
+Requires-Dist: moreau[cuda13]>=0.4.1; extra == 'cuda13'
+Provides-Extra: extras
+Requires-Dist: marimo>=0.24.2; extra == 'extras'
+Requires-Dist: matplotlib>=3.11.2; extra == 'extras'
+Requires-Dist: pandas>=3.0.6; extra == 'extras'
+Requires-Dist: planetmagfields>=1.7.0; extra == 'extras'
+Requires-Dist: seaborn>=0.13.2; extra == 'extras'
+Requires-Dist: wandb>=0.30.0; extra == 'extras'
+Description-Content-Type: text/markdown
+
+# Learning-Based Adaptive MPC for Magnetorquer-Only Satellite Attitude Control
+<img width="867" height="492" alt="hero_white_background" src="https://github.com/user-attachments/assets/659210b3-bb81-446d-8291-b291ccb81eb3" />
+
+## Prerequisites
+1. [Docker](https://docker-curriculum.com/)
+2. GPU Access
+3. [Moreau](https://docs.moreau.so/index.html#) (requires license)
+
+## Startup
+1. In workstation's bash shell, in desired directory, clone this repo
+```bash
+git clone git@github.com:JHU-ACEL/uranus-mpc.git
+```
+2. In the directory you just cloned, follow the steps in `docker/README.md` to modify the `Dockerfile` to include your Moreau License Key and build the docker image.
+4. To verify that the image built properly, run
+```bash
+docker images
+```
+4. To run a container, run
+```bash
+./docker/docker_run.sh <image_repository:tag> <PORT> # example: "./docker/docker_run.sh uranus-mpc:v1 8888"
+```
+5. To start marimo, run 
+```bash
+marimo edit --headless --host 0.0.0.0 --port=$PORT
+```
+6. If you are accessing a workstation with a GPU virtually, you will need to port forward to your local machine. In your local command line, run
+```bash
+ssh -L localhost:${PORT}:localhost:${PORT} -N -f ${UID}@{remote} # ex: "ssh -L localhost:8888:localhost:8888 -N -f pschwa24@enceladus.wse.jhu.edu"
+```
+7. On your browser on your local machine you can go to URL given in marimo (e.g. https://0.0.0.0:8888) and copy and paste the access token to access marimo.
+
+## Examples
+- `Spacecraft_Control.py` demonstrates the full functionality of the library and controller with online adaptation.
+- `B_field_Learning.py` shows the training pipeline for training the magnetic field models offline.
+
+## Working with the library
+The library is set-up to be fully compatible with JAX, and thus everything is parallelized and JIT-compatible. It makes use of [Equinox](https://docs.kidger.site/equinox/all-of-equinox/) to make this easy and set-up with Object-Oriented Programming.
+
+### Base functions
+First, set the `dynamics` class. The spacecraft attitude dynamics with magnetorquer-attitude control is initialized as follows:
+```bash
+spacecraft_dynamics = SpacecraftDynamics(mag_model=model_s,planet=planet)
+```
+
+Most of the primary functionality is built into the `TrajectoryGenerator` class, which is initialized by providing it with dynamics and a time-step.
+```bash
+system = TrajectoryGenerator(dynamics=spacecraft_dynamics, dt=dt)
+```
+
+After initializing the class, you can generate a trajectory. If the only thing provided is the initial state, key, and number of steps, the initial states are just propagated according to the dynamics. 
+```bash
+trajectory, controls = system.generate_trajectory(initial_state=initial_state, key=key, num_steps=num_steps)
+```
+
+Alternatively, you could provide a fixed control sequence that it will propagate.
+```bash
+trajectory, controls = system.generate_trajectory(initial_state=initial_state, key=key, num_steps=num_steps, control_sequence=control_sequence)
+```
+
+If your dynamics require and external dynamics parameter (e.g. magnetic field values), then those must also be provided to the `external_dynamics_params` field. Additionally, a random disturbance force/torque is applied at each step if a standard deviation for it is provided.
+```bash
+trajectory, controls = system.generate_trajectory(initial_state=init_state, target_state=target_state, key=key1, external_dynamics_params=b_field_values, control_sequence=control_sequence, num_steps=num_steps, noise_std=noise_std_dyn)
+```
+
+You can also propagate controlled trajectories. First, initialize the controllers
+```bash
+qp_controller = QPController(system, Q_qp, Qf_qp, R_qp, state_limits, control_limits, horizon_qp) # high_level_controller
+fb_controller = FeedbackController(system,Q_fb,Qf_fb,R_fb,control_limits) # low_level_controller
+```
+and then provide them to the `generate_trajectory` function along with a target state and how frequently the `high_level_controller` should replan. The `high_level_controller` is called every `replan_freq` timesteps, and the `low_level_controller` is called every timestep and tracks the nominal trajectory produced by the `high_level_controller`. If no `low_level_controller` is provided, the nominal control inputs produced by the `high_level_controller` are just used at each timestep. If no `high_level_controller` is provided, then the `low_level_controlller` just works by itself trying to reach the target state provided.
+```bash
+trajectory, controls = system.generate_trajectory(initial_state=init_state, target_state=target_state, key=key, external_dynamics_params=b_field_values, high_level_controller=qp_controller, low_level_controller=fb_controller, replan_freq=replan_freq_qp, num_steps=num_steps, noise_std=noise_std_dyn)
+```
+
+A full batch of trajectories can just as easily be generated. If an arrays of `batch_size` are provided, then it will use all of them, but if you only want to change some things for each trajectory (e.g. have the same initial state for all of them, but different target states), that works too. The following example just provides one initial state, but `batch_size` different target states and b field values. The noise always changes for each trajectory. 
+```bash
+traj_batch, cntrl_batch = system.generate_trajectory_batch(initial_states=init_state, target_states=target_states, key=key, batch_size=batch_size, external_dynamics_params=b_batch, external_dynamics_params_est=None, high_level_controller=qp_controller, low_level_controller=fb_controller, replan_freq=replan_freq_qp, chunk_size=None, num_steps=num_steps_batch, noise_std=noise_std_dyn)
+```
+
+### Using learned b-field models and online-adaptation
+In order to use the learned b-field model in the controller instead of the true values, provide the cartesian coordinates in the PCI frame of the orbit to `external_dynamics_params_est`
+Online adaptation is performed if the frequency that adaptation is performed `adapt_freq` and the size of the history buffer `history` are provided. 
+```bash
+traj_adapt, cntrl_adapt = system_coarse.generate_trajectory(initial_state=init_state, target_state=target_state, key=key, external_dynamics_params=b, external_dynamics_params_est=orbit_xyz, high_level_controller=qp_controller, low_level_controller=fb_controller, replan_freq=replan_freq_qp, num_steps=num_steps, noise_std=noise_std_dyn, adapt_freq=adapt_freq_qp, history=history_qp)
+```
+
+### Plotting
+The trajectories for all states can be plotted using
+```bash
+system.plot_traj(trajectory, controls)
+```
+
+Multiple trajectories can be plotted for comparison, and you can provide additional arguments to plot the target state as a dotted line and include a legend.
+```bash
+system.plot_traj(jnp.array([traj_true, traj_full, traj_low, traj_adapt]), jnp.array([cntrl_true, cntrl_full, cntrl_low, cntrl_adapt]), target_state=target_state, legend=["Full-order", "Low-order", "Low-order w/ Adaptation"])
+```
+
+You can also plot the costs over time using 
+```bash
+system.plot_costs(jnp.array([traj_true, traj_full, traj_low, traj_adapt]), target_state,legend=["True","Full", "Dipole", "Adapt"])
+```
+
+If you want to visualize statistics for a batch of trajectories, you can use 
+```bash
+system.plot_costs(traj_batch, target_states, plot_stats=True)
+```
+and 
+```bash
+system.plot_violin_and_bar(trajs_batch, target_states, angle_threshold=angle_threshold, omega_threshold=omega_threshold, angle_stability_tol=angle_tol, omega_stability_tol=omega_tol, tail_length=tail_length, verbose=True)
+```

raddo-0.1.3/README.md

@@ -0,0 +1,111 @@
+# Learning-Based Adaptive MPC for Magnetorquer-Only Satellite Attitude Control
+<img width="867" height="492" alt="hero_white_background" src="https://github.com/user-attachments/assets/659210b3-bb81-446d-8291-b291ccb81eb3" />
+
+## Prerequisites
+1. [Docker](https://docker-curriculum.com/)
+2. GPU Access
+3. [Moreau](https://docs.moreau.so/index.html#) (requires license)
+
+## Startup
+1. In workstation's bash shell, in desired directory, clone this repo
+```bash
+git clone git@github.com:JHU-ACEL/uranus-mpc.git
+```
+2. In the directory you just cloned, follow the steps in `docker/README.md` to modify the `Dockerfile` to include your Moreau License Key and build the docker image.
+4. To verify that the image built properly, run
+```bash
+docker images
+```
+4. To run a container, run
+```bash
+./docker/docker_run.sh <image_repository:tag> <PORT> # example: "./docker/docker_run.sh uranus-mpc:v1 8888"
+```
+5. To start marimo, run 
+```bash
+marimo edit --headless --host 0.0.0.0 --port=$PORT
+```
+6. If you are accessing a workstation with a GPU virtually, you will need to port forward to your local machine. In your local command line, run
+```bash
+ssh -L localhost:${PORT}:localhost:${PORT} -N -f ${UID}@{remote} # ex: "ssh -L localhost:8888:localhost:8888 -N -f pschwa24@enceladus.wse.jhu.edu"
+```
+7. On your browser on your local machine you can go to URL given in marimo (e.g. https://0.0.0.0:8888) and copy and paste the access token to access marimo.
+
+## Examples
+- `Spacecraft_Control.py` demonstrates the full functionality of the library and controller with online adaptation.
+- `B_field_Learning.py` shows the training pipeline for training the magnetic field models offline.
+
+## Working with the library
+The library is set-up to be fully compatible with JAX, and thus everything is parallelized and JIT-compatible. It makes use of [Equinox](https://docs.kidger.site/equinox/all-of-equinox/) to make this easy and set-up with Object-Oriented Programming.
+
+### Base functions
+First, set the `dynamics` class. The spacecraft attitude dynamics with magnetorquer-attitude control is initialized as follows:
+```bash
+spacecraft_dynamics = SpacecraftDynamics(mag_model=model_s,planet=planet)
+```
+
+Most of the primary functionality is built into the `TrajectoryGenerator` class, which is initialized by providing it with dynamics and a time-step.
+```bash
+system = TrajectoryGenerator(dynamics=spacecraft_dynamics, dt=dt)
+```
+
+After initializing the class, you can generate a trajectory. If the only thing provided is the initial state, key, and number of steps, the initial states are just propagated according to the dynamics. 
+```bash
+trajectory, controls = system.generate_trajectory(initial_state=initial_state, key=key, num_steps=num_steps)
+```
+
+Alternatively, you could provide a fixed control sequence that it will propagate.
+```bash
+trajectory, controls = system.generate_trajectory(initial_state=initial_state, key=key, num_steps=num_steps, control_sequence=control_sequence)
+```
+
+If your dynamics require and external dynamics parameter (e.g. magnetic field values), then those must also be provided to the `external_dynamics_params` field. Additionally, a random disturbance force/torque is applied at each step if a standard deviation for it is provided.
+```bash
+trajectory, controls = system.generate_trajectory(initial_state=init_state, target_state=target_state, key=key1, external_dynamics_params=b_field_values, control_sequence=control_sequence, num_steps=num_steps, noise_std=noise_std_dyn)
+```
+
+You can also propagate controlled trajectories. First, initialize the controllers
+```bash
+qp_controller = QPController(system, Q_qp, Qf_qp, R_qp, state_limits, control_limits, horizon_qp) # high_level_controller
+fb_controller = FeedbackController(system,Q_fb,Qf_fb,R_fb,control_limits) # low_level_controller
+```
+and then provide them to the `generate_trajectory` function along with a target state and how frequently the `high_level_controller` should replan. The `high_level_controller` is called every `replan_freq` timesteps, and the `low_level_controller` is called every timestep and tracks the nominal trajectory produced by the `high_level_controller`. If no `low_level_controller` is provided, the nominal control inputs produced by the `high_level_controller` are just used at each timestep. If no `high_level_controller` is provided, then the `low_level_controlller` just works by itself trying to reach the target state provided.
+```bash
+trajectory, controls = system.generate_trajectory(initial_state=init_state, target_state=target_state, key=key, external_dynamics_params=b_field_values, high_level_controller=qp_controller, low_level_controller=fb_controller, replan_freq=replan_freq_qp, num_steps=num_steps, noise_std=noise_std_dyn)
+```
+
+A full batch of trajectories can just as easily be generated. If an arrays of `batch_size` are provided, then it will use all of them, but if you only want to change some things for each trajectory (e.g. have the same initial state for all of them, but different target states), that works too. The following example just provides one initial state, but `batch_size` different target states and b field values. The noise always changes for each trajectory. 
+```bash
+traj_batch, cntrl_batch = system.generate_trajectory_batch(initial_states=init_state, target_states=target_states, key=key, batch_size=batch_size, external_dynamics_params=b_batch, external_dynamics_params_est=None, high_level_controller=qp_controller, low_level_controller=fb_controller, replan_freq=replan_freq_qp, chunk_size=None, num_steps=num_steps_batch, noise_std=noise_std_dyn)
+```
+
+### Using learned b-field models and online-adaptation
+In order to use the learned b-field model in the controller instead of the true values, provide the cartesian coordinates in the PCI frame of the orbit to `external_dynamics_params_est`
+Online adaptation is performed if the frequency that adaptation is performed `adapt_freq` and the size of the history buffer `history` are provided. 
+```bash
+traj_adapt, cntrl_adapt = system_coarse.generate_trajectory(initial_state=init_state, target_state=target_state, key=key, external_dynamics_params=b, external_dynamics_params_est=orbit_xyz, high_level_controller=qp_controller, low_level_controller=fb_controller, replan_freq=replan_freq_qp, num_steps=num_steps, noise_std=noise_std_dyn, adapt_freq=adapt_freq_qp, history=history_qp)
+```
+
+### Plotting
+The trajectories for all states can be plotted using
+```bash
+system.plot_traj(trajectory, controls)
+```
+
+Multiple trajectories can be plotted for comparison, and you can provide additional arguments to plot the target state as a dotted line and include a legend.
+```bash
+system.plot_traj(jnp.array([traj_true, traj_full, traj_low, traj_adapt]), jnp.array([cntrl_true, cntrl_full, cntrl_low, cntrl_adapt]), target_state=target_state, legend=["Full-order", "Low-order", "Low-order w/ Adaptation"])
+```
+
+You can also plot the costs over time using 
+```bash
+system.plot_costs(jnp.array([traj_true, traj_full, traj_low, traj_adapt]), target_state,legend=["True","Full", "Dipole", "Adapt"])
+```
+
+If you want to visualize statistics for a batch of trajectories, you can use 
+```bash
+system.plot_costs(traj_batch, target_states, plot_stats=True)
+```
+and 
+```bash
+system.plot_violin_and_bar(trajs_batch, target_states, angle_threshold=angle_threshold, omega_threshold=omega_threshold, angle_stability_tol=angle_tol, omega_stability_tol=omega_tol, tail_length=tail_length, verbose=True)
+```

raddo-0.1.3/docker/README.md

@@ -0,0 +1,46 @@
+# Docker 
+
+### Modify Dockerfile
+You must modify the docker file to include your Moreau License Key.
+
+### Build the Docker Container
+Use the `docker_build.sh` script to build the Docker image:
+```bash
+./docker/docker_build.sh docker/Dockerfile uranus-mpc:v1
+```
+
+**Arguments:**
+- `dockerfile`: Path to the Dockerfile (e.g., `docker/Dockerfile`)
+- `tag_name`: Docker image tag (e.g., `uranus-mpc:v1`)
+
+The build script automatically:
+- Sets the user ID and group ID to match your current user
+- Sets the username to your current username
+- Builds the image with the specified tag
+
+### Run the Docker Container
+Use the `docker_run.sh` script to run the Docker container:
+
+```bash
+./docker/docker_run.sh uranus-mpc:v1
+```
+
+**Arguments:**
+- `tag_name`: The Docker image tag to run (e.g., `uranus-mpc:v1`)
+- `ssh_port` (optional): The port forwarding between host machine and Docker container (e.g., `8888`) 
+
+### Run Jupyter Lab (from inside the container):
+```bash
+jupyter lab --ip=0.0.0.0 --port=$PORT --no-browser --allow-root
+```
+
+### Run marimo (from inside the container):
+```bash
+marimo edit --headless --host 0.0.0.0 --port=$PORT
+```
+
+### Use SSH port forwarding to run a notebook remotely
+```bash
+ssh -L localhost:${PORT}:localhost:${PORT} -N -f ${UID}@{remote}
+```
+where `${PORT}` is the port designated when you `docker run` and `${UID}@{remote}` is the remote instance (e.g., `cauligi@enceladus.wse.jhu`).

raddo-0.1.3/pyproject.toml

@@ -0,0 +1,54 @@
+[build-system]
+requires = ["hatchling>=1.27"]
+build-backend = "hatchling.build"
+
+[project]
+name = "raddo"
+version = "0.1.3"
+dependencies = [
+   "numpy>=2.5.3",
+   "jax>=0.11.2",
+   "equinox>=0.13.8",
+   "moreau>=0.4.1",
+   "optax>=0.2.8",
+]
+requires-python = ">=3.12"
+authors = [
+  {name = "Patrick Schwartz", email = "pschwa24@jh.edu"},
+]
+maintainers = [
+  {name = "Patrick Schwartz", email = "pschwa24@jh.edu"}
+]
+description = "RApidly Adapting Data-Driven Opimization (RADDO). JAX-based trajectory generator with Moreau-based QP solver with an emphasis on adaptation and attitude dynamics"
+readme = "README.md"
+classifiers = [
+  "Programming Language :: Python"
+]
+license = "MIT"
+license-files = ["LICEN[CS]E*"]
+
+[project.urls]
+Homepage = "https://github.com/JHU-ACEL/raddo"
+
+[project.optional-dependencies]
+cuda12 = [ "jax[cuda12]>=0.11.2", "moreau[cuda12]>=0.4.1"]
+cuda13 = [ "jax[cuda13]>=0.11.2", "moreau[cuda13]>=0.4.1"]
+extras =[
+    "matplotlib>=3.11.2",
+    "pandas>=3.0.6",
+    "seaborn>=0.13.2",
+    "marimo>=0.24.2",
+    "wandb>=0.30.0",
+    "planetMagFields>=1.7.0"
+]
+
+[tool.hatch.build.targets.sdist]
+include = [
+  "src/raddo",
+  "README.md",
+  "LICENSE*",
+  "pyproject.toml",
+]
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/raddo"]

raddo-0.1.3/src/raddo/controllers/__init__.py

@@ -0,0 +1,7 @@
+from .base_controller import Controller
+from .cem_controller import CEMController
+from .feedback_controller import FeedbackController
+from .mppi_controller import MPPIController
+from .qp_controller import QPController
+from .sqp_controller import SQPController
+from .tvlqr_controller import TVLQRController