Skip to content

Commit 15c1d97

Browse files
committed
Serialization of OCP optimizers
- update version to 0.10.0a1 - check SHA256 hash for rollout functions - tets_ocp: build optimizers in setUpClass
1 parent bfb2b59 commit 15c1d97

3 files changed

Lines changed: 179 additions & 102 deletions

File tree

open-codegen/VERSION

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1 +1 @@
1-
0.9.6
1+
0.10.0a1

open-codegen/opengen/ocp/builder.py

Lines changed: 54 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,12 @@
44
import json
55
import os
66
import sys
7+
import hashlib
8+
import platform
9+
from datetime import datetime, timezone
10+
import casadi
711
import casadi.casadi as cs
12+
from importlib.metadata import version, PackageNotFoundError
813

914
from opengen.builder.optimizer_builder import OpEnOptimizerBuilder
1015
from opengen.builder.problem import Problem
@@ -172,14 +177,28 @@ def __extract_states(self, flat_solution, packed_parameters):
172177

173178
def __metadata_dict(self):
174179
"""Return the JSON-serializable optimizer manifest."""
180+
rollout_sha256 = None
181+
if self.__rollout_function is not None:
182+
rollout_sha256 = None
183+
175184
return {
185+
"manifest_version": 1,
176186
"optimizer_name": self.__optimizer_name,
177187
"target_dir": self.__target_dir,
178188
"backend_kind": self.__backend_kind,
189+
"created_at": datetime.now(timezone.utc).isoformat(),
190+
"platform": platform.platform(),
191+
"python_version": sys.version,
192+
"opengen_version": self.__safe_package_version("opengen"),
193+
"casadi_version": self.__casadi_version(),
179194
"shooting": self.__shooting.value,
180195
"nx": self.__nx,
181196
"nu": self.__nu,
182197
"horizon": self.__horizon,
198+
"decision_dimension": self.__nu * self.__horizon
199+
if self.__shooting == ShootingMethod.SINGLE
200+
else self.__horizon * (self.__nu + self.__nx),
201+
"parameter_dimension": self.__parameters.total_size(),
183202
"parameters": [
184203
{
185204
"name": definition.name,
@@ -191,8 +210,25 @@ def __metadata_dict(self):
191210
"input_slices": self.__input_slices,
192211
"state_slices": self.__state_slices,
193212
"rollout_file": "rollout.casadi" if self.__rollout_function is not None else None,
213+
"rollout_sha256": rollout_sha256,
194214
}
195215

216+
@staticmethod
217+
def __safe_package_version(package_name):
218+
"""Return the installed version of a package if available."""
219+
try:
220+
return version(package_name)
221+
except PackageNotFoundError:
222+
return None
223+
224+
@staticmethod
225+
def __casadi_version():
226+
"""Return the installed CasADi version if available."""
227+
casadi_version = getattr(casadi, "__version__", None)
228+
if casadi_version is not None:
229+
return casadi_version
230+
return GeneratedOptimizer.__safe_package_version("casadi")
231+
196232
def save(self, json_path=None):
197233
"""Save a manifest that can later recreate this optimizer.
198234
@@ -213,7 +249,10 @@ def save(self, json_path=None):
213249
metadata = self.__metadata_dict()
214250
rollout_file = metadata.get("rollout_file")
215251
if rollout_file is not None:
216-
self.__rollout_function.save(os.path.join(manifest_dir, rollout_file))
252+
rollout_path = os.path.join(manifest_dir, rollout_file)
253+
self.__rollout_function.save(rollout_path)
254+
with open(rollout_path, "rb") as fh:
255+
metadata["rollout_sha256"] = hashlib.sha256(fh.read()).hexdigest()
217256

218257
with open(json_path, "w") as fh:
219258
json.dump(metadata, fh, indent=2)
@@ -245,9 +284,16 @@ def load(cls, json_path):
245284
metadata = json.load(fh)
246285
rollout_file = metadata.get("rollout_file")
247286
if rollout_file is not None:
248-
metadata["rollout_function"] = cs.Function.load(
249-
os.path.join(os.path.dirname(json_path), rollout_file)
250-
)
287+
rollout_path = os.path.join(os.path.dirname(json_path), rollout_file)
288+
metadata["rollout_function"] = cs.Function.load(rollout_path)
289+
rollout_sha256 = metadata.get("rollout_sha256")
290+
if rollout_sha256 is not None:
291+
with open(rollout_path, "rb") as fh:
292+
current_sha256 = hashlib.sha256(fh.read()).hexdigest()
293+
if current_sha256 != rollout_sha256:
294+
raise ValueError("rollout.casadi checksum mismatch")
295+
else:
296+
raise ValueError("manifest is missing rollout_sha256")
251297
backend = cls.__load_backend(
252298
metadata["target_dir"],
253299
metadata["optimizer_name"],
@@ -489,7 +535,7 @@ def __make_backend(self, target_dir):
489535
return None, "none"
490536

491537
def build(self):
492-
"""Generate, optionally compile, and wrap the optimizer.
538+
"""Generate, compile, and wrap the optimizer.
493539
494540
:return: :class:`GeneratedOptimizer`
495541
"""
@@ -506,11 +552,13 @@ def build(self):
506552
info = builder.build()
507553
target_dir = os.path.abspath(info["paths"]["target"])
508554
backend, backend_kind = self.__make_backend(target_dir)
509-
return GeneratedOptimizer(
555+
optimizer = GeneratedOptimizer(
510556
optimizer_name=self.__metadata.optimizer_name,
511557
target_dir=target_dir,
512558
backend=backend,
513559
backend_kind=backend_kind,
514560
ocp=self.__ocp,
515561
symbolic_model=symbolic_model,
516562
)
563+
optimizer.save()
564+
return optimizer

open-codegen/test/test_ocp.py

Lines changed: 124 additions & 95 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import json
12
import os
23
import unittest
34

@@ -38,11 +39,110 @@ def get_open_local_absolute_path():
3839

3940
@classmethod
4041
def setUpOCP1(cls):
41-
pass
42+
optimizer_name = "ocp_manifest_bindings"
43+
ocp = og.ocp.OptimalControlProblem(nx=2, nu=1, horizon=3)
44+
ocp.add_parameter("x0", 2)
45+
ocp.add_parameter("xref", 2, default=[0.0, 0.0])
46+
ocp.add_parameter("q", 1, default=1)
47+
ocp.add_parameter("r", 1, default=0.1)
48+
ocp.with_dynamics(lambda x, u, param: cs.vertcat(x[0] + u[0], x[1] - u[0]))
49+
ocp.with_stage_cost(
50+
lambda x, u, param, _t:
51+
param["q"] * cs.dot(x - param["xref"], x - param["xref"]) + param["r"] * cs.dot(u, u)
52+
)
53+
ocp.with_terminal_cost(
54+
lambda x, param: cs.dot(x - param["xref"], x - param["xref"])
55+
)
56+
ocp.with_input_constraints(og.constraints.Rectangle([-1.0], [1.0]))
57+
58+
cls.ocp1_optimizer = og.ocp.OCPBuilder(
59+
ocp,
60+
metadata=og.config.OptimizerMeta().with_optimizer_name(optimizer_name),
61+
build_configuration=og.config.BuildConfiguration()
62+
.with_open_version(local_path=OcpTestCase.get_open_local_absolute_path())
63+
.with_build_directory(OcpTestCase.TEST_DIR)
64+
.with_build_mode(og.config.BuildConfiguration.DEBUG_MODE)
65+
.with_build_python_bindings(),
66+
solver_configuration=og.config.SolverConfiguration()
67+
.with_tolerance(1e-5)
68+
.with_delta_tolerance(1e-5)
69+
.with_initial_penalty(10.0)
70+
.with_penalty_weight_update_factor(1.2)
71+
.with_max_inner_iterations(500)
72+
.with_max_outer_iterations(10),
73+
).build()
74+
cls.ocp1_optimizer.save()
75+
cls.ocp1_manifest_path = os.path.join(cls.ocp1_optimizer.target_dir, "optimizer_manifest.json")
76+
cls.ocp1_rollout_path = os.path.join(cls.ocp1_optimizer.target_dir, "rollout.casadi")
4277

4378
@classmethod
4479
def setUpOCP2(cls):
45-
pass
80+
def make_problem(shooting):
81+
ocp = og.ocp.OptimalControlProblem(nx=2, nu=1, horizon=5, shooting=shooting)
82+
ocp.add_parameter("x0", 2)
83+
ocp.add_parameter("xref", 2, default=[0.0, 0.0])
84+
ocp.with_dynamics(lambda x, u, param:
85+
cs.vertcat(x[0] + u[0], x[1] - u[0]))
86+
ocp.with_stage_cost(
87+
lambda x, u, param, _t:
88+
cs.dot(x - param["xref"], x - param["xref"]) + 0.01 * cs.dot(u, u)
89+
)
90+
ocp.with_terminal_cost(
91+
lambda x, param:
92+
2.0 * cs.dot(x - param["xref"], x - param["xref"])
93+
)
94+
ocp.with_input_constraints(og.constraints.Rectangle([-0.4], [0.4]))
95+
ocp.with_path_constraint(
96+
lambda x, u, param, _t: cs.fmax(0.0, x[0] - 1.5)
97+
)
98+
if shooting == og.ocp.ShootingMethod.MULTIPLE:
99+
ocp.with_dynamics_constraints("alm")
100+
return ocp
101+
102+
solver_config = og.config.SolverConfiguration() \
103+
.with_tolerance(1e-5) \
104+
.with_delta_tolerance(1e-5) \
105+
.with_initial_penalty(10.0) \
106+
.with_penalty_weight_update_factor(1.2) \
107+
.with_max_inner_iterations(5000) \
108+
.with_max_outer_iterations(20)
109+
110+
cls.ocp2_single_optimizer = og.ocp.OCPBuilder(
111+
make_problem(og.ocp.ShootingMethod.SINGLE),
112+
metadata=og.config.OptimizerMeta().with_optimizer_name("ocp_single_tcp"),
113+
build_configuration=og.config.BuildConfiguration()
114+
.with_open_version(local_path=OcpTestCase.get_open_local_absolute_path())
115+
.with_build_directory(OcpTestCase.TEST_DIR)
116+
.with_build_mode(og.config.BuildConfiguration.DEBUG_MODE)
117+
.with_tcp_interface_config(
118+
tcp_interface_config=og.config.TcpServerConfiguration(bind_port=3391)
119+
),
120+
solver_configuration=solver_config,
121+
).build()
122+
cls.ocp2_multiple_optimizer = og.ocp.OCPBuilder(
123+
make_problem(og.ocp.ShootingMethod.MULTIPLE),
124+
metadata=og.config.OptimizerMeta().with_optimizer_name("ocp_multiple_tcp"),
125+
build_configuration=og.config.BuildConfiguration()
126+
.with_open_version(local_path=OcpTestCase.get_open_local_absolute_path())
127+
.with_build_directory(OcpTestCase.TEST_DIR)
128+
.with_build_mode(og.config.BuildConfiguration.DEBUG_MODE)
129+
.with_tcp_interface_config(
130+
tcp_interface_config=og.config.TcpServerConfiguration(bind_port=3392)
131+
),
132+
solver_configuration=solver_config,
133+
).build()
134+
135+
cls.ocp2_single_optimizer.save()
136+
cls.ocp2_multiple_optimizer.save()
137+
cls.ocp2_single_manifest_path = os.path.join(
138+
cls.ocp2_single_optimizer.target_dir, "optimizer_manifest.json")
139+
cls.ocp2_multiple_manifest_path = os.path.join(
140+
cls.ocp2_multiple_optimizer.target_dir, "optimizer_manifest.json")
141+
142+
@classmethod
143+
def setUpClass(cls):
144+
cls.setUpOCP1()
145+
cls.setUpOCP2()
46146

47147
def make_ocp(self, shooting=og.ocp.ShootingMethod.SINGLE):
48148
ocp = og.ocp.OptimalControlProblem(nx=2, nu=1, horizon=3, shooting=shooting)
@@ -324,52 +424,35 @@ def test_generated_optimizer_defaults(self):
324424
self.assertAlmostEqual(result.states[-1][1], -0.6)
325425

326426
def test_optimizer_manifest_roundtrip(self):
327-
optimizer_name = "ocp_manifest_bindings"
328-
ocp = og.ocp.OptimalControlProblem(nx=2, nu=1, horizon=3)
329-
ocp.add_parameter("x0", 2)
330-
ocp.add_parameter("xref", 2, default=[0.0, 0.0])
331-
ocp.with_dynamics(lambda x, u, param: cs.vertcat(x[0] + u[0], x[1] - u[0]))
332-
ocp.with_stage_cost(
333-
lambda x, u, param, _t: cs.dot(x - param["xref"], x - param["xref"]) + 0.1 * cs.dot(u, u)
334-
)
335-
ocp.with_terminal_cost(
336-
lambda x, param: cs.dot(x - param["xref"], x - param["xref"])
337-
)
338-
ocp.with_input_constraints(og.constraints.Rectangle([-1.0], [1.0]))
339-
340-
optimizer = og.ocp.OCPBuilder(
341-
ocp,
342-
metadata=og.config.OptimizerMeta().with_optimizer_name(optimizer_name),
343-
build_configuration=og.config.BuildConfiguration()
344-
.with_open_version(local_path=OcpTestCase.get_open_local_absolute_path())
345-
.with_build_directory(OcpTestCase.TEST_DIR)
346-
.with_build_mode(og.config.BuildConfiguration.DEBUG_MODE)
347-
.with_build_python_bindings(),
348-
solver_configuration=og.config.SolverConfiguration()
349-
.with_tolerance(1e-5)
350-
.with_delta_tolerance(1e-5)
351-
.with_initial_penalty(10.0)
352-
.with_penalty_weight_update_factor(1.2)
353-
.with_max_inner_iterations(500)
354-
.with_max_outer_iterations(10),
355-
).build()
356-
357-
optimizer.save()
358-
manifest_path = os.path.join(optimizer.target_dir, "optimizer_manifest.json")
359-
rollout_path = os.path.join(optimizer.target_dir, "rollout.casadi")
427+
manifest_path = self.ocp1_manifest_path
428+
rollout_path = self.ocp1_rollout_path
360429

361430
self.assertTrue(os.path.exists(manifest_path))
362431
self.assertTrue(os.path.exists(rollout_path))
363432

433+
with open(manifest_path, "r") as fh:
434+
manifest = json.load(fh)
435+
436+
self.assertEqual(1, manifest["manifest_version"])
437+
self.assertEqual("ocp_manifest_bindings", manifest["optimizer_name"])
438+
self.assertEqual("direct", manifest["backend_kind"])
439+
self.assertEqual("single", manifest["shooting"])
440+
self.assertEqual(2, manifest["nx"])
441+
self.assertEqual(1, manifest["nu"])
442+
self.assertEqual(3, manifest["horizon"])
443+
self.assertEqual(3, manifest["decision_dimension"])
444+
self.assertEqual(6, manifest["parameter_dimension"]) # p = (x0, xref, q, r)
445+
self.assertEqual("rollout.casadi", manifest["rollout_file"])
446+
self.assertIsNotNone(manifest["rollout_sha256"])
447+
self.assertIn("created_at", manifest)
448+
364449
loaded_optimizer = og.ocp.GeneratedOptimizer.load(manifest_path)
365450

366451
result = loaded_optimizer.solve(
367-
x0=[1.0, 0.0],
368-
xref=[0.0, 0.0],
369-
initial_guess=[0.0] * 3,
452+
x0=[1.0, 0.0]
370453
)
371454

372-
self.assertEqual(optimizer_name, loaded_optimizer.optimizer_name)
455+
self.assertEqual("ocp_manifest_bindings", loaded_optimizer.optimizer_name)
373456
self.assertEqual("direct", loaded_optimizer.backend_kind)
374457
self.assertEqual("Converged", result.exit_status)
375458
self.assertEqual(3, len(result.inputs))
@@ -440,62 +523,8 @@ def test_multiple_shooting_state_extraction(self):
440523
self.assertEqual(result.states, [[0.0, 0.0], [0.1, -0.1], [0.3, -0.3], [0.7, -0.7]])
441524

442525
def test_single_vs_multiple(self):
443-
def make_problem(shooting):
444-
ocp = og.ocp.OptimalControlProblem(nx=2, nu=1, horizon=5, shooting=shooting)
445-
ocp.add_parameter("x0", 2)
446-
ocp.add_parameter("xref", 2, default=[0.0, 0.0])
447-
ocp.with_dynamics(lambda x, u, param:
448-
cs.vertcat(x[0] + u[0], x[1] - u[0]))
449-
ocp.with_stage_cost(
450-
lambda x, u, param, _t:
451-
cs.dot(x - param["xref"], x - param["xref"]) + 0.01 * cs.dot(u, u)
452-
)
453-
ocp.with_terminal_cost(
454-
lambda x, param:
455-
2.0 * cs.dot(x - param["xref"], x - param["xref"])
456-
)
457-
ocp.with_input_constraints(og.constraints.Rectangle([-0.4], [0.4]))
458-
ocp.with_path_constraint(
459-
lambda x, u, param, _t: cs.fmax(0.0, x[0] - 1.5)
460-
)
461-
if shooting == og.ocp.ShootingMethod.MULTIPLE:
462-
ocp.with_dynamics_constraints("alm")
463-
return ocp
464-
465-
solver_config = og.config.SolverConfiguration() \
466-
.with_tolerance(1e-5) \
467-
.with_delta_tolerance(1e-5) \
468-
.with_initial_penalty(10.0) \
469-
.with_penalty_weight_update_factor(1.2) \
470-
.with_max_inner_iterations(5000) \
471-
.with_max_outer_iterations(20)
472-
473-
ocp_single = make_problem(og.ocp.ShootingMethod.SINGLE)
474-
single_optimizer = og.ocp.OCPBuilder(
475-
ocp_single,
476-
metadata=og.config.OptimizerMeta().with_optimizer_name("ocp_single_tcp"),
477-
build_configuration=og.config.BuildConfiguration()
478-
.with_open_version(local_path=OcpTestCase.get_open_local_absolute_path())
479-
.with_build_directory(OcpTestCase.TEST_DIR)
480-
.with_build_mode(og.config.BuildConfiguration.DEBUG_MODE)
481-
.with_tcp_interface_config(
482-
tcp_interface_config=og.config.TcpServerConfiguration(bind_port=3391)
483-
),
484-
solver_configuration=solver_config,
485-
).build()
486-
ocp_multiple = make_problem(og.ocp.ShootingMethod.MULTIPLE)
487-
multiple_optimizer = og.ocp.OCPBuilder(
488-
ocp_multiple,
489-
metadata=og.config.OptimizerMeta().with_optimizer_name("ocp_multiple_tcp"),
490-
build_configuration=og.config.BuildConfiguration()
491-
.with_open_version(local_path=OcpTestCase.get_open_local_absolute_path())
492-
.with_build_directory(OcpTestCase.TEST_DIR)
493-
.with_build_mode(og.config.BuildConfiguration.DEBUG_MODE)
494-
.with_tcp_interface_config(
495-
tcp_interface_config=og.config.TcpServerConfiguration(bind_port=3392)
496-
),
497-
solver_configuration=solver_config,
498-
).build()
526+
single_optimizer = og.ocp.GeneratedOptimizer.load(self.ocp2_single_manifest_path)
527+
multiple_optimizer = og.ocp.GeneratedOptimizer.load(self.ocp2_multiple_manifest_path)
499528

500529
try:
501530
x0 = [1.0, -1.0]

0 commit comments

Comments
 (0)