alphaville / optimization-engine

Nonconvex embedded optimization: code generation for fast real-time optimization + ROS support
https://alphaville.github.io/optimization-engine/
Other
499 stars 53 forks source link

RFC: Automatically generating python bindings #215

Closed davidrusu closed 3 years ago

davidrusu commented 3 years ago

Hello! I've been experimenting with generating python bindings to the generated rust code, this would be an alternative to the existing TCP interface.

I've gone through the process of handcoding some bindings to one of my solvers using pyo3, I've attached the code below, it think it was relatively straight forward.

In practice it looks like this:

# from the cargo project with the generated pyo3 api:
> cargo build --release
> mv target/release/libopen_agent.so open_agent.so # hopefully we can automate this
> python
>>> import open_agent
>>> solver = open_agent.build_solver()
>>> result = solver.run(params)
>>> result.solution
[0.2, 0.1....]

Since I'm running the solver on the same machine as the robot (in simulation), I'm seeing more reliable and performant behavior, we get to skip the serialization, and TCP stack and keep everything within the same process boundary. (I periodically saw retries while the robot was waiting for the solver's TCP server to spin up. For long running simulations this can eat up valuable compute time).

Below is my experimental code generating the python bindings, far from perfect but I wanted to get comments on it asap, treat it as a sketch of a solution.

Please let me know your thoughts. Are you open to adding support for automatically generating python bindings? If so, I'd be happy to continue developing this idea, but if you'd prefer to take it over, that's perfectly fine too.

use optimization_engine::alm::*;

use pyo3::prelude::*;
use pyo3::wrap_pyfunction;

use trajectory_optimizer_v176::*;

#[pymodule]
fn open_agent(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(build_solver, m)?)?;
    m.add_class::<OptimizerSolution>()?;
    m.add_class::<Solver>()?;
    Ok(())
}

#[pyfunction]
fn build_solver() -> PyResult<Solver> {
    let cache = initialize_solver();
    Ok(Solver { cache })
}

/// Solution and solution status of optimizer
#[pyclass]
struct OptimizerSolution {
    #[pyo3(get)]
    exit_status: String,
    #[pyo3(get)]
    num_outer_iterations: usize,
    #[pyo3(get)]
    num_inner_iterations: usize,
    #[pyo3(get)]
    last_problem_norm_fpr: f64,
    #[pyo3(get)]
    delta_y_norm_over_c: f64,
    #[pyo3(get)]
    f2_norm: f64,
    #[pyo3(get)]
    solve_time_ms: f64,
    #[pyo3(get)]
    penalty: f64,
    #[pyo3(get)]
    solution: Vec<f64>,
    #[pyo3(get)]
    lagrange_multipliers: Vec<f64>,
    #[pyo3(get)]
    cost: f64,
}

#[pyclass]
struct Solver {
    cache: AlmCache,
}

#[pymethods]
impl Solver {
    /// Run solver
    fn run(
        &mut self,
        parameter: Vec<f64>,
        initial_guess: Option<Vec<f64>>,
        initial_lagrange_multipliers: Option<Vec<f64>>,
        initial_penalty: Option<f64>,
    ) -> PyResult<Option<OptimizerSolution>> {
        let mut u = [0.0; TRAJECTORY_OPTIMIZER_V176_NUM_DECISION_VARIABLES];

        // ----------------------------------------------------
        // Set initial value
        // ----------------------------------------------------
        if let Some(u0) = initial_guess {
            if u0.len() != TRAJECTORY_OPTIMIZER_V176_NUM_DECISION_VARIABLES {
                println!(
                    "1600 -> Initial guess has incompatible dimensions: {} != {}",
                    u0.len(),
                    TRAJECTORY_OPTIMIZER_V176_NUM_DECISION_VARIABLES
                );
                return Ok(None);
            }
            u.copy_from_slice(&u0);
        }

        // ----------------------------------------------------
        // Check lagrange multipliers
        // ----------------------------------------------------
        if let Some(y0) = &initial_lagrange_multipliers {
            if y0.len() != TRAJECTORY_OPTIMIZER_V176_N1 {
                println!(
                    "1700 -> wrong dimension of Langrange multipliers: {} != {}",
                    y0.len(),
                    TRAJECTORY_OPTIMIZER_V176_N1
                );
                return Ok(None);
            }
        }

        // ----------------------------------------------------
        // Check parameter
        // ----------------------------------------------------
        if parameter.len() != TRAJECTORY_OPTIMIZER_V176_NUM_PARAMETERS {
            println!(
                "3003 -> wrong number of parameters: {} != {}",
                parameter.len(),
                TRAJECTORY_OPTIMIZER_V176_NUM_PARAMETERS
            );
            return Ok(None);
        }

        // ----------------------------------------------------
        // Run solver
        // ----------------------------------------------------
        let solver_status = solve(
            &parameter,
            &mut self.cache,
            &mut u,
            &initial_lagrange_multipliers,
            &initial_penalty,
        );

        match solver_status {
            Ok(status) => Ok(Some(OptimizerSolution {
                exit_status: format!("{:?}", status.exit_status()),
                num_outer_iterations: status.num_outer_iterations(),
                num_inner_iterations: status.num_inner_iterations(),
                last_problem_norm_fpr: status.last_problem_norm_fpr(),
                delta_y_norm_over_c: status.delta_y_norm_over_c(),
                f2_norm: status.f2_norm(),
                penalty: status.penalty(),
                lagrange_multipliers: status.lagrange_multipliers().clone().unwrap_or_default(),
                solve_time_ms: (status.solve_time().as_nanos() as f64) / 1e6,
                solution: u.to_vec(),
                cost: status.cost(),
            })),
            Err(_) => {
                println!("2000 -> Problem solution failed (solver error)");
                Ok(None)
            }
        }
    }
}
alphaville commented 3 years ago

Hi @davidrusu. Thank you very much for this RFC. I think this is a great idea and can certainly be automated using Jinja2 templates for code generation.

One possible approach is that when .with_build_c_bindings() is used, OpEn will generate C bindings and a separate Rust project with pyo3 (like the code you wrote above).

I created a new branch called feature/pyo3-interface . Would you be willing to create an example with the above code and put it in that branch? A hand-coded example should be fine for now, just to experiment with it a bit.

davidrusu commented 3 years ago

Sure thing, today is looking a little busy, I'll put something up tomorrow :+1:

alphaville commented 3 years ago

@davidrusu, I looked into the #217, tested it on Ubuntu by running main.py and everything works fine. I introduced a test in opengen/test/test.py, which is currently running on Travis CI (it will take a few minutes).

I found this in the documentation of pyo3:

While developing, you can symlink (or copy) and rename the shared library from the target folder: On MacOS, rename libstring_sum.dylib to string_sum.so, on Windows libstring_sum.dll to string_sum.pyd, and on Linux libstring_sum.so to string_sum.so. Then open a Python shell in the same folder and you'll be able to import string_sum.

It seems that we need to modify the code for Windows and Mac OSX**.

In the last line of OpEnOptimizerBuilder.__build_python_bindings you copy the generated so file to the current directory. My concern is that the current working directory may not be in the part.

For example in test.py I had to do the following:

import sys
import os
sys.path.insert(1, os.getcwd())  # include the CWD into the path!

I wonder whether there is a better way to go about it. For instance, in your example, we could move the .so file into my_optimizers/halfspace_optimizer/ and then include that into the path.

** ~as expected, the test failed on Mac OS X on Travis CI.~ Fixed!

alphaville commented 3 years ago

A first version of this feature is available in version 0.6.4 (read the documentation). There will need to be a few updates in future versions.

davidrusu commented 3 years ago

Been running the 0.6.4 version all morning and it's been great. Episode resets are now much quicker and reliable and the solver performances is much more reliable (no more retries :+1:)