snibo13/pendulum.py

## pendulum.py
import casadi as cas
import numpy as np
import time
import matplotlib.pyplot as plt

import mujoco
import mujoco.viewer
import numpy as np
import matplotlib.pyplot as plt
import time


m = 1
l = 1
g = 9.81

dt = 0.01


def accel(x, u):
    I = m * l**2
    tau_g = m * g * l * cas.sin(x[0])
    net_torque = u - tau_g
    return net_torque / I


def rk4_a(x, u):
    k1 = accel(x, u)
    k2 = accel(x + k1 * dt / 2, u)
    k3 = accel(x + k2 * dt / 2, u)
    k4 = accel(x + k3 * dt, u)
    return (k1 + 2 * k2 + 2 * k3 + k4) * dt / 6


def rk4_v(x, u):
    k1 = x[1]
    k2 = x[1] + rk4_a(x, u) * dt / 2
    k3 = x[1] + rk4_a(x, u) * dt / 2
    k4 = x[1] + rk4_a(x, u) * dt
    return (k1 + 2 * k2 + 2 * k3 + k4) * dt / 6


Q = np.diag([100, 100])
Q_ = 100 * Q

opti = cas.Opti()
T = 5
N = int(T / dt)


x = opti.variable(2, N + 1)
y = opti.variable(1, N)

J = 0
target = np.array([np.pi, 0])
x0 = np.array([0, 0])
print(x0)
opti.subject_to(x[:, 0] == x0)
for i in range(N):
    J += (x[:, i] - target).T @ Q @ (x[:, i] - target)
    J += y[:, i].T @ y[:, i]
    opti.subject_to(x[1, i + 1] == x[1, i] + rk4_a(x[:, i], y[:, i]))
    opti.subject_to(x[0, i + 1] == x[0, i] + rk4_v(x[:, i], y[:, i]))
    opti.subject_to(opti.bounded(-30, y[:, i], 30))


J += (x[:, -1] - target).T @ Q_ @ (x[:, -1] - target)


opti.minimize(J)

opti.solver("ipopt", {"ipopt.print_level": 0})
# try:
sol = opti.solve()
x = sol.value(x)
y = sol.value(y)
plt.subplot(4, 1, 1)
plt.plot(x[0, :], label="Position")
plt.legend()
plt.subplot(4, 1, 2)
plt.plot(x[1, :], label="Velocity")
plt.legend()
plt.subplot(4, 1, 3)
acceleration = np.zeros(N)
for i in range(N):
    acceleration[i] = accel(x[:, i], y[i])

plt.plot(acceleration, label="Acceleration")
plt.legend()
plt.subplot(4, 1, 4)
plt.plot(y, label="Applied Force")
plt.legend()
plt.show()
# except Exception as e:
#     print(e)

print("Final position: ", x[:, -1])
print("Target: ", target)
print("Max applied force: ", np.max(y))
print("Final error: ", np.linalg.norm(x[:, -1] - target) ** 2)


model_path = "pendulum.xml"
model = mujoco.MjModel.from_xml_path(model_path)
data = mujoco.MjData(model)
thetas = np.zeros(N)
velocity = np.zeros(N)
with mujoco.viewer.launch_passive(model, data) as viewer:
    viewer.cam.lookat = [0, 0, 0]
    viewer.cam.distance = 15
    viewer.cam.azimuth = -90
    viewer.cam.elevation = 0
    for i in range(N):
        data.ctrl[0] = y[i]
        mujoco.mj_step(model, data)
        thetas[i] = data.qpos[0]
        velocity[i] = data.qvel[0]
        viewer.sync()
        time.sleep(0.01)
    viewer.close()

plt.subplot(2, 1, 1)
plt.plot(thetas, label="Theta")
plt.plot(x[0, :], label="Optimal")
plt.legend()
plt.subplot(2, 1, 2)
plt.plot(velocity, label="Velocity")
plt.plot(x[1, :], label="Optimal")
plt.legend()
plt.show()

## pendulum.xml
<mujoco model="mass">
    <option gravity="0 0 -9.81" timestep="0.01" integrator="RK4">
        <flag constraint="disable" energy="enable"/>
    </option>
    <compiler autolimits="true" />
    <asset>
        <texture type="skybox" builtin="gradient" rgb1="0 0 0" rgb2=".6 .8 1" width="256" height="256"/>
        <texture name="groundtex" type="2d" builtin="checker" rgb1="1 1 1" rgb2="0.3 0.3 .8" width="512" height="512" mark="edge" markrgb="1 1 1"/>
        <material name="groundplane" texture="groundtex" texrepeat="1 1" texuniform="true"/>
        <material name="lower" rgba="1 0 0 1"/>
        <material name="upper" rgba="0 0.7 0 1"/>
        <material name="foot" rgba="0 0 0.7 1"/>
        <material name="mass" rgba="0 0 0 1" />
    </asset>
    <worldbody>
        <light pos="0 0 2" dir="-1 -1 -1" target="pendulum" diffuse="1 1 1"/>
        <geom type="plane" size="20 20 .1" material="groundplane"/>

        <body name="pendulum" pos="0 0 2">
            <joint name="pin" type="hinge" pos="0 0 0" axis="0 1 0" damping="0"/>
            <geom type="capsule" size="0.1" fromto="0 0 0 0 0 -1" material="lower" mass="0"/>
            <geom type="sphere" pos="0 0 -1" size="0.1" material="upper" mass="1"/>
        </body>
    </worldbody>
    <actuator>
        <motor name="joint" joint="pin" forcerange="-30 30" ctrlrange="-30 30"/>
    </actuator>


</mujoco>
	import casadi as cas
	import numpy as np
	import time
	import matplotlib.pyplot as plt

	import mujoco
	import mujoco.viewer
	import numpy as np
	import matplotlib.pyplot as plt
	import time


	m = 1
	l = 1
	g = 9.81

	dt = 0.01


	def accel(x, u):
	I = m * l**2
	tau_g = m * g * l * cas.sin(x[0])
	net_torque = u - tau_g
	return net_torque / I


	def rk4_a(x, u):
	k1 = accel(x, u)
	k2 = accel(x + k1 * dt / 2, u)
	k3 = accel(x + k2 * dt / 2, u)
	k4 = accel(x + k3 * dt, u)
	return (k1 + 2 * k2 + 2 * k3 + k4) * dt / 6


	def rk4_v(x, u):
	k1 = x[1]
	k2 = x[1] + rk4_a(x, u) * dt / 2
	k3 = x[1] + rk4_a(x, u) * dt / 2
	k4 = x[1] + rk4_a(x, u) * dt
	return (k1 + 2 * k2 + 2 * k3 + k4) * dt / 6


	Q = np.diag([100, 100])
	Q_ = 100 * Q

	opti = cas.Opti()
	T = 5
	N = int(T / dt)


	x = opti.variable(2, N + 1)
	y = opti.variable(1, N)

	J = 0
	target = np.array([np.pi, 0])
	x0 = np.array([0, 0])
	print(x0)
	opti.subject_to(x[:, 0] == x0)
	for i in range(N):
	J += (x[:, i] - target).T @ Q @ (x[:, i] - target)
	J += y[:, i].T @ y[:, i]
	opti.subject_to(x[1, i + 1] == x[1, i] + rk4_a(x[:, i], y[:, i]))
	opti.subject_to(x[0, i + 1] == x[0, i] + rk4_v(x[:, i], y[:, i]))
	opti.subject_to(opti.bounded(-30, y[:, i], 30))


	J += (x[:, -1] - target).T @ Q_ @ (x[:, -1] - target)


	opti.minimize(J)

	opti.solver("ipopt", {"ipopt.print_level": 0})
	# try:
	sol = opti.solve()
	x = sol.value(x)
	y = sol.value(y)
	plt.subplot(4, 1, 1)
	plt.plot(x[0, :], label="Position")
	plt.legend()
	plt.subplot(4, 1, 2)
	plt.plot(x[1, :], label="Velocity")
	plt.legend()
	plt.subplot(4, 1, 3)
	acceleration = np.zeros(N)
	for i in range(N):
	acceleration[i] = accel(x[:, i], y[i])

	plt.plot(acceleration, label="Acceleration")
	plt.legend()
	plt.subplot(4, 1, 4)
	plt.plot(y, label="Applied Force")
	plt.legend()
	plt.show()
	# except Exception as e:
	# print(e)

	print("Final position: ", x[:, -1])
	print("Target: ", target)
	print("Max applied force: ", np.max(y))
	print("Final error: ", np.linalg.norm(x[:, -1] - target) ** 2)


	model_path = "pendulum.xml"
	model = mujoco.MjModel.from_xml_path(model_path)
	data = mujoco.MjData(model)
	thetas = np.zeros(N)
	velocity = np.zeros(N)
	with mujoco.viewer.launch_passive(model, data) as viewer:
	viewer.cam.lookat = [0, 0, 0]
	viewer.cam.distance = 15
	viewer.cam.azimuth = -90
	viewer.cam.elevation = 0
	for i in range(N):
	data.ctrl[0] = y[i]
	mujoco.mj_step(model, data)
	thetas[i] = data.qpos[0]
	velocity[i] = data.qvel[0]
	viewer.sync()
	time.sleep(0.01)
	viewer.close()

	plt.subplot(2, 1, 1)
	plt.plot(thetas, label="Theta")
	plt.plot(x[0, :], label="Optimal")
	plt.legend()
	plt.subplot(2, 1, 2)
	plt.plot(velocity, label="Velocity")
	plt.plot(x[1, :], label="Optimal")
	plt.legend()
	plt.show()
	<mujoco model="mass">
	<option gravity="0 0 -9.81" timestep="0.01" integrator="RK4">
	<flag constraint="disable" energy="enable"/>
	</option>
	<compiler autolimits="true" />
	<asset>
	<texture type="skybox" builtin="gradient" rgb1="0 0 0" rgb2=".6 .8 1" width="256" height="256"/>
	<texture name="groundtex" type="2d" builtin="checker" rgb1="1 1 1" rgb2="0.3 0.3 .8" width="512" height="512" mark="edge" markrgb="1 1 1"/>
	<material name="groundplane" texture="groundtex" texrepeat="1 1" texuniform="true"/>
	<material name="lower" rgba="1 0 0 1"/>
	<material name="upper" rgba="0 0.7 0 1"/>
	<material name="foot" rgba="0 0 0.7 1"/>
	<material name="mass" rgba="0 0 0 1" />
	</asset>
	<worldbody>
	<light pos="0 0 2" dir="-1 -1 -1" target="pendulum" diffuse="1 1 1"/>
	<geom type="plane" size="20 20 .1" material="groundplane"/>

	<body name="pendulum" pos="0 0 2">
	<joint name="pin" type="hinge" pos="0 0 0" axis="0 1 0" damping="0"/>
	<geom type="capsule" size="0.1" fromto="0 0 0 0 0 -1" material="lower" mass="0"/>
	<geom type="sphere" pos="0 0 -1" size="0.1" material="upper" mass="1"/>
	</body>
	</worldbody>
	<actuator>
	<motor name="joint" joint="pin" forcerange="-30 30" ctrlrange="-30 30"/>
	</actuator>


	</mujoco>