In [ ]:
# imports from standard libs
import mcts
import matplotlib.pyplot as plt
import networkx as nx
# chrono imports
import pychrono as chrono

from control_optimisation import create_grab_criterion_fun, create_traj_fun, get_object_to_grasp
from rostok.graph_grammar.node import GraphGrammar
from rostok.trajectory_optimizer.control_optimizer import ConfigRewardFunction, ControlOptimizer
from rostok.criterion.flags_simualtions import FlagMaxTime, FlagSlipout, FlagNotContact
import rule_extention
import rostok.graph_generators.graph_environment as env

This tutorial shows how to start the search for the optimal design of the mechanism built with predetermined sets of nodes and rules.
First step is to import rules and and nodes constructed in another tutorial

In [ ]:
# Create extension rule vocabulary
rule_vocabul, node_features = rule_extention.init_extension_rules()

Configure the optimizer as shown in optimization tutorial. The result of this step is the configured control_optimizer object of ControlOptimizer class

In [ ]:
# Create condig optimizing control

GAIT = 2.5
WEIGHT = [5, 0, 1, 9]

cfg = ConfigRewardFunction()
cfg.bound = (0, 10)
cfg.iters = 2
cfg.sim_config = {"Set_G_acc": chrono.ChVectorD(0, 0, 0)}
cfg.time_step = 0.001
cfg.time_sim = 3
cfg.flags = [FlagMaxTime(3), FlagNotContact(1), FlagSlipout(1, 0.25)]
"""Wraps function call"""

criterion_callback = create_grab_criterion_fun(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

control_optimizer = ControlOptimizer(cfg)

Here we use Monte Carlo Tree Search (MCTS) to search the space of possible graphs. This algorithm build several random solutions that can be obtained using rules from the current point and then make a step (apply rule to change current point) using the calculated rewards. Set two parameters of searching:

  1. limit on amount of solutions built to perform a step
  2. max number of non-terminal rules for a solution

Create an instance of mcts class using the iteration_limit parameter to initialize the searcher.

In [ ]:
# Hyperparameters mctss
iteration_limit = 10
max_numbers_rules = 12
# Initialize MCTS
searcher = mcts.mcts(iterationLimit=iteration_limit)

The starting point is a graph with only ROOT node. One can apply rules at this stage to make another starting state. Create an instance of GraphVocabularyEnvironment class using the starting state, your rule vocabulary and max number of rules. Add optimizer to the environment

In [ ]:
graph = GraphGrammar()

# Create graph envirenments for algorithm 
graph_env = env.GraphVocabularyEnvironment(graph, rule_vocabul, max_numbers_rules)

graph_env.set_control_optimizer(control_optimizer)

Start the searching loop. At each step searcher check several solutions and returns a rule that leads to the most promising results. This rule is applied by step function. The algorithm stops when after a step each node of the graph is terminal, and hence we obtained a solution physical solution.

In [ ]:
# Run  algorithm
iter = 0
finish = False
while not finish:
    action = searcher.search(initialState=graph_env)
    finish, final_graph, opt_trajectory = graph_env.step(action, False)
    iter += 1
    print(
        f"number iteration: {iter}, counter actions: {graph_env.counter_action}, reward: {graph_env.reward}"
    )

Plot the resulting solution

In [ ]:
def plot_graph(graph: GraphGrammar):
    plt.figure()
    nx.draw_networkx(graph,
                     pos=nx.kamada_kawai_layout(graph, dim=2),
                     node_size=800,
                     labels={n: graph.nodes[n]["Node"].label for n in graph})
    plt.show()

plot_graph(final_graph)