File size: 4,266 Bytes
b791e4c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 | """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))
# Sum the feedback before adding the offsets, as in the trained policy.
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))
|