| """Export the skating, rolling and learned landing actions as one ONNX policy.""" |
|
|
| import learn |
| import numpy as np |
| import onnx |
| from onnx import TensorProto, helper, numpy_helper |
|
|
|
|
| def export(weights, destination): |
| nodes, constants = [], [] |
|
|
| def constant(name, value, dtype=np.float32): |
| constants.append(numpy_helper.from_array(np.asarray(value, dtype=dtype), name)) |
| return name |
|
|
| def node(kind, name, *inputs, **attributes): |
| nodes.append(helper.make_node(kind, inputs, [name], name=name, **attributes)) |
| return name |
|
|
| def column(name, ix): |
| return node( |
| "Gather", name, "obs", constant(name + "_idx", [ix], np.int64), axis=1 |
| ) |
|
|
| def scaled_clip(value, factor, name): |
| scaled = node( |
| "Mul", name + "_scaled", value, constant(name + "_factor", factor) |
| ) |
| return node("Clip", name, scaled, neg_one, one) |
|
|
| zero = constant("zero", 0) |
| one = constant("one", 1) |
| neg_one = constant("neg_one", -1) |
| cue = column("cue", 60) |
| phase = [ |
| node("Greater", "after_" + str(i), cue, constant("boundary_" + str(i), i + 0.5)) |
| for i in range(5) |
| ] |
| models = {} |
| for name in learn.BASE_MODELS: |
| file = learn.model_file(name) |
| source = onnx.compose.add_prefix(onnx.load(str(file)), name + "/") |
| scale = 1.12 if name == "roll" else 1.0 |
| divisor = np.ones(61, np.float32) |
| divisor[34:48] = scale |
| mask = np.ones(61, np.float32) |
| mask[48 if name == "roll" else 49 :] = 0 |
| prepared = node( |
| "Div", name + "_obs", "obs", constant(name + "_divisor", divisor) |
| ) |
| node( |
| "Mul", source.graph.input[0].name, prepared, constant(name + "_mask", mask) |
| ) |
| nodes.extend(source.graph.node) |
| constants.extend(source.graph.initializer) |
| models[name] = node( |
| "Mul", |
| name + "_action", |
| source.graph.output[0].name, |
| constant(name + "_gain", scale), |
| ) |
| base = node("Where", "recovery_base", phase[3], models["rise"], models["roll"]) |
| arr = np.asarray(weights, np.float32) |
| biases = arr[:24].reshape(4, 6) |
| bias = constant("phase2", biases[0]) |
| for n in range(1, 4): |
| bias = node( |
| "Where", |
| "bias" + str(n), |
| phase[n + 1], |
| constant("phase" + str(n + 2), biases[n]), |
| bias, |
| ) |
| pitch_gravity = column("pitch_gravity", 3) |
| roll_gravity = column("roll_gravity", 4) |
| pitch_rate = scaled_clip(column("pitch_rate_body", 1), 0.14, "pitch_rate") |
| roll_rate = scaled_clip(column("roll_rate_body", 0), 0.15, "roll_rate") |
|
|
| def vector_term(obs, name, arr): |
| return node("Mul", name, obs, constant(name + "_weights", arr)) |
|
|
| |
| pitch_terms = node( |
| "Add", |
| "pitch_feedback", |
| vector_term(pitch_gravity, "pitch_gravity_term", arr[24:27]), |
| vector_term(pitch_rate, "pitch_rate_term", arr[27:30]), |
| ) |
| roll_terms = node( |
| "Add", |
| "roll_feedback", |
| vector_term(roll_gravity, "roll_gravity_term", arr[30]), |
| vector_term(roll_rate, "roll_rate_term", arr[31]), |
| ) |
| zeros = node("Mul", "zeros", pitch_rate, zero) |
| feedback = node("Concat", "feedback", pitch_terms, zeros, zeros, roll_terms, axis=1) |
| activated = node("Tanh", "bounded", node("Add", "raw", bias, feedback)) |
| adjustment = node( |
| "MatMul", "joint_adjustment", activated, constant("joint_map", learn.MAP) |
| ) |
| corrected = node("Add", "corrected", base, adjustment) |
| roll_or_recovery = node("Where", "trick", phase[1], corrected, base) |
| node("Where", "actions", phase[0], roll_or_recovery, models["skate"]) |
| graph = helper.make_graph( |
| nodes, |
| "skate_land_skate", |
| [helper.make_tensor_value_info("obs", TensorProto.FLOAT, [1, 61])], |
| [helper.make_tensor_value_info("actions", TensorProto.FLOAT, [1, 14])], |
| constants, |
| ) |
| model = helper.make_model( |
| graph, ir_version=8, opset_imports=[helper.make_opsetid("", 18)] |
| ) |
| onnx.checker.check_model(model, full_check=True) |
| onnx.save(model, str(destination)) |
|
|
