In [ ]:
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

In [ ]:
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.

In [ ]:
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.

In [ ]:
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:

  • Size of the space of simulation
  • iters
  • Properties of a simulation (here it is gravity parameters)
  • Time step for simulation (seconds)
  • Max time of the simulation (seconds)
  • Set the flags for simulation (here only max_time)
In [ ]:
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.

In [ ]:
"""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

In [ ]:
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