import example_vocabulary
import pychrono as chrono
import rostok.criterion.criterion_calc as criterion
from rostok.block_builder.basic_node_block import SimpleBody
from rostok.criterion.flags_simualtions import FlagMaxTime, FlagNotContact
from rostok.graph_grammar.node import BlockWrapper, GraphGrammar, Node
from rostok.trajectory_optimizer.control_optimizer import (ConfigRewardFunction, ControlOptimizer)
from rostok.trajectory_optimizer.trajectory_generator import create_torque_traj_from_x
from rostok.virtual_experiment.simulation_step import SimOut
from rostok.block_builder.node_render import DefaultChronoMaterial, FrameTransform, ChronoBodyEnv
The optimization of joints trajectories is performed by the object of ControlOptimizer
class. But one need to configure the optimizer using ConfigRewardFunction
class.
At first we will create utility functions that we will use in optimization. First function will return object to grasp. More about configuring an object see in sym_step tutorial
def get_object_to_grasp():
matich = DefaultChronoMaterial()
matich.Friction = 0.65
matich.DampingF = 0.65
obj = BlockWrapper(ChronoBodyEnv,
shape=SimpleBody.CYLINDER,
material=matich,
pos=FrameTransform([0, 1, 0], [0, -0.048, 0.706, 0.706]))
return obj
We need a function that returns reward for robot and simulation output.
def grab_crtitrion(sim_output: dict[int, SimOut], grab_robot, node_feature: list[list[Node]], gait,
weight):
j_nodes = criterion.nodes_division(grab_robot, node_feature[1])
b_nodes = criterion.nodes_division(grab_robot, node_feature[0])
rb_nodes = criterion.sort_left_right(grab_robot, node_feature[3], node_feature[0])
lb_nodes = criterion.sort_left_right(grab_robot, node_feature[2], node_feature[0])
return criterion.criterion_calc(sim_output, b_nodes, j_nodes, rb_nodes, lb_nodes, weight, gait)
def create_grab_criterion_fun(node_features, gait, weight):
def fun(sim_output, grab_robot):
return grab_crtitrion(sim_output, grab_robot, node_features, gait, weight)
return fun
We need a function that takes graph and the list of scalar parameters for joints and returns the trajectories for joints.
def create_traj_fun(stop_time: float, time_step: float):
def fun(graph: GraphGrammar, x: list[float]):
return create_torque_traj_from_x(graph, x, stop_time, time_step)
return fun
To setup the optimization one need to configure the reward function. Create an instance of ConfigRewardFunction
class. Set the parameters:
GAIT = 2.5
WEIGHT = [5, 0, 1, 5]
cfg = ConfigRewardFunction()
cfg.bound = (-5, 5)
cfg.iters = 2
cfg.sim_config = {"Set_G_acc": chrono.ChVectorD(0, 0, 0)}
cfg.time_step = 0.005
cfg.time_sim = 2
cfg.flags = [FlagMaxTime(cfg.time_sim)]
The configuration includes functions that we created above, now we need to add them to the configuration.
"""Wraps function call"""
criterion_callback = create_grab_criterion_fun(example_vocabulary.NODE_FEATURES, GAIT, WEIGHT)
traj_generator_fun = create_traj_fun(cfg.time_sim, cfg.time_step)
cfg.criterion_callback = criterion_callback
cfg.get_rgab_object_callback = get_object_to_grasp
cfg.params_to_timesiries_callback = traj_generator_fun
Crete an instance of ControlOptimizer
using your configuration object and pass the graph of the mechanism you want to get optimization for to the start_optimisation
function of the ControlOptimizer
object. The result is the array of optimal coefficients for the joints and the corresponding best reward
control_optimizer = ControlOptimizer(cfg)
graph = example_vocabulary.get_terminal_graph_three_finger()
res = control_optimizer.start_optimisation(graph)
print(res)
(-4.482628750730302, array([0. , 0. , 0. , 0. , 0. , 3.33333333]))
Here we obtained the best reward for one mechanism. In the next tutorial we will use the optimization configuration to search the space of possible mechanisms trying to find the best morphology possible for the given set of nodes and rules