118 lines
3.7 KiB
Python
118 lines
3.7 KiB
Python
import pytest
|
|
import numpy as np
|
|
|
|
from PythonLinearNonlinearControl.models.cartpole import CartPoleModel
|
|
from PythonLinearNonlinearControl.configs.cartpole \
|
|
import CartPoleConfigModule
|
|
|
|
class TestCartPoleModel():
|
|
"""
|
|
"""
|
|
def test_step(self):
|
|
config = CartPoleConfigModule()
|
|
cartpole_model = CartPoleModel(config)
|
|
|
|
curr_x = np.ones(4)
|
|
curr_x[2] = np.pi / 6.
|
|
|
|
us = np.ones((1, 1))
|
|
|
|
next_x = cartpole_model.predict_traj(curr_x, us)
|
|
|
|
d_x0 = curr_x[1]
|
|
d_x1 = (1. + config.MP * np.sin(np.pi / 6.) \
|
|
* (config.L * (1.**2) \
|
|
+ config.G * np.cos(np.pi / 6.))) \
|
|
/ (config.MC + config.MP * np.sin(np.pi / 6.)**2)
|
|
d_x2 = curr_x[3]
|
|
d_x3 = (-1. * np.cos(np.pi / 6.) \
|
|
- config.MP * config.L * (1.**2) \
|
|
* np.cos(np.pi / 6.) * np.sin(np.pi / 6.) \
|
|
- (config.MP + config.MC) * config.G \
|
|
* np.sin(np.pi / 6.)) \
|
|
/ (config.L \
|
|
* (config.MC \
|
|
+ config.MP * np.sin(np.pi / 6.)**2))
|
|
|
|
expected = np.array([d_x0, d_x1, d_x2, d_x3]) * config.DT \
|
|
+ curr_x
|
|
|
|
expected = np.stack((curr_x, expected), axis=0)
|
|
|
|
assert next_x == pytest.approx(expected, abs=1e-5)
|
|
|
|
def test_predict_traj(self):
|
|
config = CartPoleConfigModule()
|
|
cartpole_model = CartPoleModel(config)
|
|
|
|
curr_x = np.ones(config.STATE_SIZE)
|
|
curr_x[-1] = np.pi / 6.
|
|
u = np.ones((1, config.INPUT_SIZE))
|
|
|
|
pred_xs = cartpole_model.predict_traj(curr_x, u)
|
|
|
|
u = np.tile(u, (2, 1, 1))
|
|
pred_xs_alltogether = cartpole_model.predict_traj(curr_x, u)[0]
|
|
|
|
assert pred_xs_alltogether == pytest.approx(pred_xs)
|
|
|
|
def test_gradient_state(self):
|
|
|
|
config = CartPoleConfigModule()
|
|
cartpole_model = CartPoleModel(config)
|
|
|
|
xs = np.ones((1, config.STATE_SIZE)) \
|
|
* np.random.rand(1, config.STATE_SIZE)
|
|
xs[0, -1] = np.pi / 6.
|
|
us = np.ones((1, config.INPUT_SIZE))
|
|
|
|
grad = cartpole_model.calc_f_x(xs, us, config.DT)
|
|
|
|
# expected cost
|
|
expected_grad = np.zeros((1, config.STATE_SIZE, config.STATE_SIZE))
|
|
eps = 1e-4
|
|
|
|
for i in range(config.STATE_SIZE):
|
|
tmp_x = xs.copy()
|
|
tmp_x[0, i] = xs[0, i] + eps
|
|
forward = \
|
|
cartpole_model.predict_next_state(tmp_x[0], us[0])
|
|
tmp_x = xs.copy()
|
|
tmp_x[0, i] = xs[0, i] - eps
|
|
backward = \
|
|
cartpole_model.predict_next_state(tmp_x[0], us[0])
|
|
|
|
expected_grad[0, :, i] = (forward - backward) / (2. * eps)
|
|
|
|
assert grad == pytest.approx(expected_grad)
|
|
|
|
def test_gradient_input(self):
|
|
|
|
config = CartPoleConfigModule()
|
|
cartpole_model = CartPoleModel(config)
|
|
|
|
xs = np.ones((1, config.STATE_SIZE)) \
|
|
* np.random.rand(1, config.STATE_SIZE)
|
|
xs[0, -1] = np.pi / 6.
|
|
us = np.ones((1, config.INPUT_SIZE))
|
|
|
|
grad = cartpole_model.calc_f_u(xs, us, config.DT)
|
|
|
|
# expected cost
|
|
expected_grad = np.zeros((1, config.STATE_SIZE, config.INPUT_SIZE))
|
|
eps = 1e-4
|
|
|
|
for i in range(config.INPUT_SIZE):
|
|
tmp_u = us.copy()
|
|
tmp_u[0, i] = us[0, i] + eps
|
|
forward = \
|
|
cartpole_model.predict_next_state(xs[0], tmp_u[0])
|
|
tmp_u = us.copy()
|
|
tmp_u[0, i] = us[0, i] - eps
|
|
backward = \
|
|
cartpole_model.predict_next_state(xs[0], tmp_u[0])
|
|
|
|
expected_grad[0, :, i] = (forward - backward) / (2. * eps)
|
|
|
|
assert grad == pytest.approx(expected_grad)
|