optimization.py

import mujoco
import mujoco_viewer
import matplotlib.pyplot as plt
import numpy as np
import os
import lxml
from lxml import etree
import mujoco.viewer
import time
import csv
from actuators import actuators

visualisationActive = False

#X_mech params
Xn = 0.08 # distance between both hips
l1 = 0.1 # length of first parts of both legs
l2 = 0.1 # lenght of second parts of both legs

totalNumber = len(actuators)-1
currentActuator = 0

KP = 200
KD = 5

itr = 1
step = 0.02
desiredTime = 30
delay = 1


csv_filename = "jump_results.csv"
with open(csv_filename, mode='w', newline='', encoding='utf-8') as csvfile:
    writer = csv.writer(csvfile)
    writer.writerow(["Итерация","Привод","Xn","l1","l2","","Начальное положение тела","Максимальная высота прыжка","Максимальное время полёта"])

while currentActuator < 16.1:
    while Xn < 0.21:
        while l1 < 0.41:
            while l2 < 0.41:

                modelName = list(actuators.keys())[currentActuator]
                XControl = actuators[modelName]

                print(f"Iteration №{itr}")
                print(f"Actuator model: {modelName}")
                print(f"l2 = {l2}")
                print(f"l1 = {l1}")
                print(f"Xn = {Xn}")
                print("============")

                f1 = "modelStart.xml"
                f2 = "model.xml"

                def swap_par(tree, element_type, element_name, attribute_name, new_value):
                    element = tree.find(f'.//{element_type}[@name="{element_name}"]')
                    element.set(attribute_name, new_value)

                tree = etree.parse(f1)

                # Xn parsing of hips
                swap_par(tree, 'body', 'leftHip', 'pos', f"-{Xn} -0.1 0")
                swap_par(tree, 'body', 'rightHip', 'pos', f"{Xn} -0.1 0")

                # l1 and l2 parsing of the left leg
                swap_par(tree, 'body', 'leftL1', 'pos', f"0 -{l1} 0.03")
                swap_par(tree, 'geom', 'leftL1', 'size', f"0.05 {l1} 0.03")
                swap_par(tree, 'body', 'leftKnee', 'pos', f"0 -{l1} 0")
                swap_par(tree, 'body', 'leftL2', 'pos', f"0 -{l2} -0.06")
                swap_par(tree, 'geom', 'leftL2', 'size', f"0.03 {l2} 0.05")
                swap_par(tree, 'body', 'endeffector', 'pos', f"0 -{l2} 0")

                # l1 and l2 parsing of the right leg
                swap_par(tree, 'body', 'rightL1', 'pos', f"0 -{l1} 0.03")
                swap_par(tree, 'geom', 'rightL1', 'size', f"0.05 {l1} 0.03")
                swap_par(tree, 'body', 'rightKnee', 'pos', f"0 -{l1} 0")
                swap_par(tree, 'body', 'rightL2', 'pos', f"0 -{l2} -0.06")
                swap_par(tree, 'geom', 'rightL2', 'size',  f"0.03 {l2} 0.05")
                swap_par(tree, 'body', 'linking', 'pos', f"0 -{l2} 0")

                #All X Control params for both motors
                swap_par(tree, 'joint', 'hipLeft', 'damping', f"{XControl["rated_torque"]/((XControl["rated_speed"]*2*np.pi)/60)}")
                swap_par(tree, 'joint', 'hipRight', 'damping', f"{XControl["rated_torque"]/((XControl["rated_speed"]*2*np.pi)/60)}")

                swap_par(tree, 'motor', 'leftHipMotor', 'ctrlrange', f"-{XControl["rated_torque"]} {XControl["rated_torque"]}")
                swap_par(tree, 'motor', 'rightHipMotor', 'ctrlrange', f"-{XControl["rated_torque"]} {XControl["rated_torque"]}")
                swap_par(tree, 'motor', 'leftHipMotor', 'gear', f"{XControl["reduction_ratio"]}")
                swap_par(tree, 'motor', 'rightHipMotor', 'gear', f"{XControl["reduction_ratio"]}")

                tree.write(f2, pretty_print=True, xml_declaration=True, encoding='UTF-8')

                model = mujoco.MjModel.from_xml_path(f2)
                data = mujoco.MjData(model)

                def set_torque(mj_data, KP, KV, theta):
                    data.ctrl[0] = KP * (-mj_data.qpos[0] + theta) + KV * (0 - mj_data.qvel[0])
                    data.ctrl[1] = KP * (-mj_data.qpos[0] + theta) + KV * (0 - mj_data.qvel[0])

                SIMEND = desiredTime
                TIMESTEP = 0.001
                STEP_NUM = int(SIMEND / TIMESTEP)
                timeseries = np.linspace(0, SIMEND, STEP_NUM)

                T = 2
                FREQ = 1/T
                AMP = np.deg2rad(-90)

                theta_des = 5*AMP * np.sin(FREQ * timeseries)

                tStart = 0
                tFinish = 0
                tFlight = 0
                jumpHeight = 0
                bodyStart = 0

                if visualisationActive:
                    viewer = mujoco_viewer.MujocoViewer(model, data, title="jumpingMechanism", width=1920, height=1080)
                else:
                    viewer = None

                for i in range(STEP_NUM):
                    if visualisationActive and not viewer.is_alive:
                        break

                    if data.time < delay:
                        data.ctrl[:] = 0
                    else:
                        set_torque(data, KP, KD, theta_des[i])

                    position_EE = data.site_xpos[3]
                    position_Body = data.site_xpos[0]

                    if bodyStart == 0:
                        if round(position_EE[2], 3) == 0:
                            bodyStart = round(position_Body[2], 3)

                    # Max Jump
                    if (round(position_Body[2], 3) > bodyStart and data.time > delay and round(position_Body[2], 3) > jumpHeight):
                        jumpHeight = round(position_Body[2], 3)

                    # Max Time
                    if (abs(round(position_EE[2], 3)) > 0.03 and round(position_Body[2], 3) > bodyStart and data.time > delay):
                        if tStart == 0:
                            tStart = data.time
                    if (round(position_Body[2], 3) <= bodyStart and data.time > delay and tStart != 0 and abs(round(position_EE[2], 3)) <= 0.03):
                        if tFinish == 0:
                            tFinish = data.time
                    if tFinish != 0 and tStart != 0:
                        if (tFinish - tStart) > tFlight:
                            tFlight = tFinish - tStart
                        tStart = 0
                        tFinish = 0

                    mujoco.mj_step(model, data)

                    if visualisationActive:
                        viewer.render()

                if visualisationActive:
                    viewer.close()

                print(f'Начальное положение тела: {bodyStart}')
                print(f'Самый высокий прыжок: {jumpHeight}')
                print(f'Максимальное время полёта: {tFlight}')
                print("============")

                with open(csv_filename, mode='a', newline='', encoding='utf-8') as csvfile:
                    writer = csv.writer(csvfile)
                    writer.writerow([itr,modelName,Xn,l1,l2,"",bodyStart,jumpHeight,round(tFlight, 3)])

                itr += 1
                l2 = round(l2 + step, 2)

            l2 = 0.1
            l1 = round(l1 + step, 2)

        l1 = 0.1                         
        Xn = round(Xn + step*2, 2)
    
    Xn = 0.08
    currentActuator += 1
print("Simulation is over!")