From ab715ad16af7849e5d7c182bb1a81f2a35e88618 Mon Sep 17 00:00:00 2001 From: Filippo Airaldi Date: Mon, 27 May 2024 13:31:45 +0200 Subject: [PATCH] fixing tests 2 --- examples/bayesopt.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/examples/bayesopt.py b/examples/bayesopt.py index b1f325b..6d0d513 100644 --- a/examples/bayesopt.py +++ b/examples/bayesopt.py @@ -36,7 +36,6 @@ from gymnasium import Env, ObservationWrapper from gymnasium.spaces import Box from gymnasium.wrappers import TimeLimit, TransformReward -from joblib import Parallel from scipy.stats.qmc import LatinHypercube from mpcrl import ( @@ -246,12 +245,11 @@ def get_cstr_mpc( env: CstrEnv, horizon: int, multistarts: int, n_jobs: int ) -> Mpc[cs.SX]: """Returns an MPC controller for the given CSTR env.""" - # create a custom parallel object for the multistart NLP - parallel = Parallel(n_jobs=n_jobs, return_as="generator") - parallel.__enter__() # not strictly necessary, but good practice - - # create NLP instance and MPC wrapper - nlp = ParallelMultistartNlp[cs.SX]("SX", starts=multistarts, parallel=parallel) + nlp = ParallelMultistartNlp[cs.SX]( + "SX", + starts=multistarts, + parallel_kwargs={"n_jobs": n_jobs, "return_as": "generator"}, + ) mpc = Mpc[cs.SX](nlp, horizon) # variables (state, action)