def test_to_from_euler(self):
"""2-way conversion starting from quaternions."""
np.random.seed(0)
angles_euler = np.pi * np.random.rand(100, 3)
conventions_euler = ["xzx", "xyx", "yxy", "yzy", "zyz", "zxz"]
# For Tait-Bryan angles the second angle must be between -pi/2 and pi/2
angles_tb = angles_euler.copy()
angles_tb[:, 1] -= np.pi / 2
conventions_tb = ["xzy", "xyz", "yxz", "yzx", "zyx", "zxy"]
axis_types = ["extrinsic", "intrinsic"]
for convention in conventions_euler:
for axis_type in axis_types:
out = rowan.to_euler(
rowan.from_euler(
angles_euler[..., 0],
angles_euler[..., 1],
angles_euler[..., 2],
convention,
axis_type,
),
convention,
axis_type,
)
self.assertTrue(
np.all(
np.logical_or(
np.isclose(out - angles_euler, 0),
np.isclose(out + angles_euler, 0),
)),
msg="Failed for convention {}, axis type {}".format(
convention, axis_type),
)
for convention in conventions_tb:
for axis_type in axis_types:
out = rowan.to_euler(
rowan.from_euler(
angles_tb[..., 0],
angles_tb[..., 1],
angles_tb[..., 2],
convention,
axis_type,
),
convention,
axis_type,
)
self.assertTrue(
np.all(
np.logical_or(
np.isclose(out - angles_tb, 0),
np.isclose(out + angles_tb, 0),
)),
msg="Failed for convention {}, axis type {}".format(
convention, axis_type),
)
def test_simple_translation():
box = freud.box.Box(10, 10, 10)
motion = TrackedMotion(box, np.zeros((1, 3)),
rowan.from_euler([0], [0], [0]))
for i in range(100):
motion.add(np.ones((1, 3)) * (i + 1), rowan.from_euler([0], [0], [0]))
assert np.isclose(np.linalg.norm(motion.delta_translation),
np.sqrt(3) * 100)
def test_simple_orientation():
box = freud.box.Box(10, 10, 10)
motion = TrackedMotion(box, np.zeros((1, 3)),
rowan.from_euler([0], [0], [0]))
for i in range(10):
motion.add(np.zeros((1, 3)),
rowan.from_euler([np.pi / 4 * (i + 1)], [0], [0]))
assert np.isclose(np.linalg.norm(motion.delta_rotation), np.pi / 4 * 10)
def test_properties(angle):
assume(np.cos(angle))
neighbourlist = np.array([[1], [0]])
orientation = np.array(
[rowan.from_euler(0, 0, 0),
rowan.from_euler(angle, 0, 0)])
result = _orientational_order(neighbourlist, orientation)
assert np.allclose(result, np.square(np.cos(angle)))
def test_translation(translations):
box = freud.box.Box(10, 10, 10)
motion = TrackedMotion(box, np.zeros((1, 3)),
rowan.from_euler([0], [0], [0]))
positions = np.cumsum(translations, axis=0)
for position in positions:
motion.add(position, rowan.from_euler([0], [0], [0]))
print(positions[-1])
print(motion.delta_translation)
assert np.allclose(motion.delta_translation, positions[-1], atol=1e-5)
def test_orientation(rotations):
box = freud.box.Box(10, 10, 10)
motion = TrackedMotion(box, np.zeros((1, 3)),
rowan.from_euler([0], [0], [0]))
orientations = np.cumsum(rotations)
for orientation in orientations:
motion.add(np.zeros((1, 3)), rowan.from_euler([orientation], [0], [0]))
print(sum(rotations))
print(motion.delta_rotation.sum())
assert np.isclose(motion.delta_rotation[0, 0],
np.sum(rotations),
atol=1e-8)
def test_from_to_euler(self):
"""2-way conversion starting from Euler angles."""
np.random.seed(0)
quats = rowan.normalize(np.random.rand(25, 4))
conventions = [
"xzx",
"xyx",
"yxy",
"yzy",
"zyz",
"zxz",
"xzy",
"xyz",
"yxz",
"yzx",
"zyx",
"zxy",
]
axis_types = ["extrinsic", "intrinsic"]
for convention in conventions:
for axis_type in axis_types:
euler = rowan.to_euler(quats, convention, axis_type)
out = rowan.from_euler(euler[..., 0], euler[..., 1],
euler[..., 2], convention, axis_type)
self.assertTrue(
np.all(
np.logical_or(np.isclose(out - quats, 0),
np.isclose(out + quats, 0))),
msg="Failed for convention {}, axis type {}".format(
convention, axis_type),
)
def z2quaternion(theta: np.ndarray) -> np.ndarray:
"""Convert a rotation about the z axis to a quaternion.
This is a helper for 2D simulations, taking the rotation of a particle about the z axis and
converting it to a quaternion. The input angle `theta` is assumed to be in radians.
"""
return rowan.from_euler(theta, 0, 0).astype(np.float32)
def test_perpendicular():
neighbourlist = np.array([[1], [0]])
orientation = np.zeros((2, 4))
orientation[:, 0] = 1
orientation[0, :] = rowan.from_euler(np.pi / 2, 0, 0)
result = _orientational_order(neighbourlist, orientation)
assert np.allclose(result, 0)
def test_large_rotation():
orientation = util.zero_quaternion(10)
dyn = dynamics.Dynamics(0, np.ones(3), np.zeros((10, 3)), orientation)
for i in range(10):
dyn.add(
np.zeros((10, 3)),
rowan.from_euler(
np.ones(10) * np.pi / 4 * i, np.zeros(10), np.zeros(10)),
)
assert np.allclose(np.linalg.norm(dyn.delta_rotation, axis=1),
np.pi / 4 * 9)
def _calc_shift(orientations: np.ndarray, molecule: Molecule) -> np.ndarray:
"""A function to calculate the shift of large particle positions to COM positions.
The positions defined in the Trimer classes are for the center of the 'large' particle since
this was the output from the packing algorithm. This offsets these values to be on the center of
mass, taking into account the orientation of the molecule.
Args:
orientations: Orientation of particles in degrees
molecule: The molecule for which the shift is occurring.
"""
pos_shift = molecule.get_relative_positions()[0]
orient_quat = rowan.from_euler(np.deg2rad(orientations), 0, 0)
return rowan.rotate(orient_quat, pos_shift)
def test_struct(translation, rotation):
translation = np.append(translation, np.zeros((1, 1)), axis=1)
init_pos = np.zeros((1, 3))
init_orient = util.zero_quaternion(1)
dyn = dynamics.Dynamics(0,
np.ones(3) * 10,
init_pos,
init_orient,
wave_number=2.9,
molecule=Trimer())
orientation = rowan.from_euler(rotation[:, 0], rotation[:, 1], rotation[:,
2])
dyn.add(translation, orientation)
assert np.allclose(dyn.delta_rotation, rotation, atol=2e-7)
assert np.allclose(dyn.delta_translation, translation, atol=8e-7)
def reset(self, initial_state=None):
if initial_state is None:
self.s = np.empty(self.n)
# position and velocity
limits = np.array([0.5, 0.5, 0.5, 1, 1, 1, 0, 0, 0, 0, 12, 12, 12])
self.s[0:6] = np.random.uniform(-limits[0:6], limits[0:6], 6)
# rotation
rpy = np.radians(
np.random.uniform(-self.rpy_limit, self.rpy_limit, 3))
q = rowan.from_euler(rpy[0], rpy[1], rpy[2], 'xyz')
self.s[6:10] = q
# angular velocity
self.s[10:13] = np.random.uniform(-limits[10:13], limits[10:13], 3)
else:
self.s = initial_state
self.time_step = 0
return np.array(self.s)
def compute_struct_relax(self) -> float:
if self.distance is None:
raise ValueError(
"The wave number is required for the structural relaxation.")
quat_rot = rowan.from_euler(
self.delta_rotation[:, 2],
self.delta_rotation[:, 1],
self.delta_rotation[:, 0],
)
final_pos = molecule2particles(self.delta_translation, quat_rot,
self.mol_vector)
init_pos = molecule2particles(
np.zeros_like(self.delta_translation),
zero_quaternion(self.num_particles),
self.mol_vector,
)
try:
return np.mean(
np.linalg.norm(final_pos - init_pos, axis=1) < self.distance)
except FloatingPointError:
return np.nan
def reset(self, initial_state=None):
if initial_state is None:
self.s = np.empty(self.n)
# position and velocity
self.s[0:6] = np.random.uniform(-self.limits[0:6],
self.limits[0:6], 6)
# rotation
if self.rpy_limit is None:
q = rowan.random.random_sample()
else:
rpy = np.radians(
np.random.uniform(-self.rpy_limit, self.rpy_limit, 3))
q = rowan.from_euler(rpy[0], rpy[1], rpy[2], 'xyz')
if q[0] < 0:
q = -q
self.s[6:10] = q
# angular velocity
self.s[10:13] = np.random.uniform(-self.limits[10:13],
self.limits[10:13], 3)
else:
self.s = initial_state
self.time_step = 0
return np.array(self.s)
def unit_quaternion_Z():
return angle().map(lambda z: rowan.from_euler(z, 0, 0))
def read(self):
"Read the frame data from the stream."
self.stream.seek(self.start)
i = line = None
def _assert(assertion):
assert i is not None
assert line is not None
if not assertion:
raise ParserError("Failed to read line #{}: {}.".format(
i, line))
monotype = False
raw_frame = _RawFrameData()
raw_frame.view_rotation = None
for i, line in enumerate(self.stream):
if _is_comment(line):
continue
if line.startswith('#'):
if line.startswith('#[data]'):
_assert(raw_frame.data is None)
raw_frame.data_keys, raw_frame.data, j = \
self._read_data_section(line, self.stream)
i += j
else:
raise ParserError(line)
else:
tokens = line.rstrip().split()
if not tokens:
continue # empty line
elif tokens[0] in TOKENS_SKIP:
continue # skip these lines
if tokens[0] == 'eof':
logger.debug("Reached end of frame.")
break
elif tokens[0] == 'def':
definition, data, end = line.strip().split('"')
_assert(len(end) == 0)
name = definition.split()[1]
type_shape = self._parse_shape_definition(data)
if name not in raw_frame.types:
raw_frame.types.append(name)
raw_frame.type_shapes.append(type_shape)
else:
typeid = raw_frame.type_shapes.index(name)
raw_frame.type_shapes[typeid] = type_shape
warnings.warn(
"Redefinition of type '{}'.".format(name))
elif tokens[0] == 'shape': # monotype
data = line.strip().split('"')[1]
raw_frame.types.append(self.default_type)
type_shape = self._parse_shape_definition(data)
raw_frame.type_shapes.append(type_shape)
_assert(len(raw_frame.type_shapes) == 1)
monotype = True
elif tokens[0] in ('boxMatrix', 'box'):
if len(tokens) == 10:
raw_frame.box = np.array(
[self._num(v) for v in tokens[1:]]).reshape((3, 3))
elif len(tokens) == 4:
raw_frame.box = np.array([[self._num(tokens[1]), 0, 0],
[0,
self._num(tokens[2]), 0],
[0, 0,
self._num(tokens[3])]
]).reshape((3, 3))
elif tokens[0] == 'rotation':
euler_angles = np.array([float(t) for t in tokens[1:]])
euler_angles *= np.pi / 180
raw_frame.view_rotation = rowan.from_euler(
*euler_angles, axis_type='extrinsic', convention='xyz')
else:
# assume we are reading positions now
if not monotype:
name = tokens[0]
if name not in raw_frame.types:
raw_frame.types.append(name)
type_shape = self._parse_shape_definition(' '.join(
tokens[:3]))
raw_frame.type_shapes.append(type_shape)
else:
name = self.default_type
typeid = raw_frame.types.index(name)
type_shape = raw_frame.type_shapes[typeid]
if len(tokens) == 7 and isinstance(type_shape, ArrowShape):
xyz = tokens[-6:-3]
quat = tokens[-3:] + [0]
elif len(tokens) >= 7:
xyz = tokens[-7:-4]
quat = tokens[-4:]
elif len(tokens) >= 3:
xyz = tokens[-3:]
quat = None
else:
raise ParserError(line)
raw_frame.typeid.append(typeid)
raw_frame.position.append([self._num(v) for v in xyz])
if quat is None:
raw_frame.orientation.append(quat)
else:
raw_frame.orientation.append(
[self._num(v) for v in quat])
# Perform inverse rotation to recover original coordinates
if raw_frame.view_rotation is not None:
pos = rowan.rotate(rowan.inverse(raw_frame.view_rotation),
raw_frame.position)
else:
pos = np.asarray(raw_frame.position)
# If all the z coordinates are close to zero, set box dimension to 2
if np.allclose(pos[:, 2], 0.0, atol=1e-7):
raw_frame.box_dimensions = 2
# If no valid orientations have been added, the array should be empty
if all([quat is None for quat in raw_frame.orientation]):
raw_frame.orientation = []
else:
# Replace values of None with an identity quaternion
for i in range(len(raw_frame.orientation)):
if raw_frame.orientation[i] is None:
raw_frame.orientation[i] = [1, 0, 0, 0]
return raw_frame
def test_from_euler(self):
"""Convert Euler angles to quaternions"""
alpha, beta, gamma = [0, 0, 0]
self.assertTrue(
np.all(
rowan.from_euler(alpha, beta, gamma) == np.array([1, 0, 0, 0
])))
alpha, beta, gamma = [np.pi / 2, np.pi / 2, 0]
self.assertTrue(
np.allclose(
rowan.from_euler(alpha, beta, gamma, 'zyx', 'intrinsic'),
np.array([0.5, -0.5, 0.5, 0.5])))
# Confirm broadcasting works from different Euler angles
alpha, beta, gamma = [[0, np.pi / 2], [0, np.pi / 2], 0]
self.assertTrue(
np.allclose(rowan.from_euler(alpha, beta, gamma),
np.array([[1, 0, 0, 0], [0.5, -0.5, 0.5, 0.5]])))
alpha = [[0, np.pi / 2], [0, np.pi / 2]]
beta = [0, np.pi / 2]
gamma = 0
self.assertTrue(
np.allclose(
rowan.from_euler(alpha, beta, gamma),
np.array([[[1, 0, 0, 0], [0.5, -0.5, 0.5, 0.5]],
[[1, 0, 0, 0], [0.5, -0.5, 0.5, 0.5]]])))
# More complicated test, checks 2d arrays and more complex angles
alpha = euler_angles[:, 0]
beta = euler_angles[:, 1]
gamma = euler_angles[:, 2]
self.assertTrue(
np.allclose(
rowan.from_euler(alpha, beta, gamma, 'zyz', 'intrinsic'),
euler_quaternions))
# Ensure proper errors are raised
with self.assertRaises(ValueError):
rowan.from_euler(alpha, beta, gamma, 'foo', 'intrinsic')
with self.assertRaises(ValueError):
rowan.from_euler(alpha, beta, gamma, 'foo', 'extrinsic')
with self.assertRaises(ValueError):
rowan.from_euler(alpha, beta, gamma, 'zyz', 'bar')
def test_zero_beta(self):
"""Check cases where beta is 0."""
# Since the Euler calculations are all done using matrices, it's easier
# to construct the test cases by directly using matrices as well. We
# assume gamma is 0 since, due to gimbal lock, only either alpha+gamma
# or alpha-gamma is a relevant parameter, and we just scan the other
# possible values. The actual function is defined such that gamma will
# always be zero in those cases. We define the matrices using lambda
# functions to support sweeping a range of values for alpha and beta,
# specifically to test cases where signs flip e.g. cos(0) vs cos(pi).
# These sign flips lead to changes in the rotation angles that must be
# tested.
mats_euler_intrinsic = [
('xzx', 'intrinsic', lambda alpha, beta:
[[np.cos(beta), 0, 0],
[0, np.cos(beta) * np.cos(alpha), -np.sin(alpha)],
[0, np.cos(beta) * np.sin(alpha),
np.cos(alpha)]]),
('xyx', 'intrinsic', lambda alpha, beta:
[[np.cos(beta), 0, 0],
[0, np.cos(alpha), -np.cos(beta) * np.sin(alpha)],
[0, np.sin(alpha),
np.cos(beta) * np.cos(alpha)]]),
('yxy', 'intrinsic', lambda alpha, beta:
[[np.cos(alpha), 0,
np.cos(beta) * np.sin(alpha)], [0, np.cos(beta), 0],
[-np.sin(alpha), 0,
np.cos(beta) * np.cos(alpha)]]),
('yzy', 'intrinsic', lambda alpha, beta:
[[np.cos(beta) * np.cos(alpha), 0,
np.sin(alpha)], [0, np.cos(beta), 0],
[-np.cos(beta) * np.sin(alpha), 0,
np.cos(alpha)]]),
('zyz', 'intrinsic', lambda alpha, beta: [[
np.cos(beta) * np.cos(alpha), -np.sin(alpha), 0
], [np.cos(beta) * np.sin(alpha),
np.cos(beta), 0], [0, 0, np.cos(beta)]]),
('zxz', 'intrinsic', lambda alpha, beta:
[[np.cos(alpha), -np.cos(beta) * np.sin(alpha), 0],
[np.sin(alpha), np.cos(beta) * np.cos(beta), 0],
[0, 0, np.cos(beta)]]),
]
mats_tb_intrinsic = [
('xzy', 'intrinsic', lambda alpha, beta: [[0, -np.sin(
beta), 0], [np.sin(beta) * np.cos(alpha), 0, -np.sin(
alpha)], [np.sin(beta) * np.sin(alpha), 0,
np.cos(alpha)]]),
('xyz', 'intrinsic',
lambda alpha, beta: [[0, 0, np.sin(
beta)], [np.sin(beta) * np.sin(alpha),
np.cos(alpha), 0
], [-np.sin(beta) * np.cos(alpha),
np.sin(alpha), 0]]),
('yxz', 'intrinsic', lambda alpha, beta:
[[np.cos(alpha), np.sin(beta) * np.sin(alpha), 0],
[0, 0, -np.sin(beta)],
[-np.sin(alpha), np.sin(beta) * np.cos(alpha), 0]]),
('yzx', 'intrinsic', lambda alpha, beta:
[[0, -np.sin(beta) * np.cos(alpha),
np.sin(alpha)], [np.sin(beta), 0, 0],
[0, np.sin(beta) * np.sin(alpha),
np.cos(alpha)]]),
('zyx', 'intrinsic', lambda alpha, beta: [[
0, -np.sin(alpha),
np.sin(beta) * np.cos(alpha)
], [0, np.cos(alpha),
np.sin(beta) * np.sin(alpha)], [-np.sin(beta), 0, 0]]),
('zxy', 'intrinsic', lambda alpha, beta: [[
np.cos(alpha), 0,
np.sin(beta) * np.sin(alpha)
], [np.sin(alpha), 0, -np.sin(beta) * np.cos(alpha)], [0, -1, 0]]),
]
# Extrinsic rotations can be tested identically to intrinsic rotations
# in the case of proper Euler angles.
mats_euler_extrinsic = [(m[0], 'extrinsic', m[2])
for m in mats_euler_intrinsic]
# For Tait-Bryan angles, extrinsic rotations axis order must be
# reversed (since axes 1 and 3 are not identical), but more
# importantly, due to the sum/difference of alpha and gamma that
# arises, we need to test the negative of alpha to catch the dangerous
# cases. In practice we get the same results since we're sweeping alpha
# values in the tests below, but it's useful to set this up precisely.
mats_tb_extrinsic = [(m[0][::-1], 'extrinsic', lambda alpha, beta: m[2]
(-alpha, beta)) for m in mats_tb_intrinsic]
# Since angle representations may not be unique, checking that
# quaternions are equal may not work. Instead we perform rotations and
# check that they are identical. For simplicity, we rotate the
# simplest vector with all 3 components (otherwise tests won't catch
# the problem because there's no component to rotate).
test_vector = [1, 1, 1]
mats_intrinsic = (mats_euler_intrinsic, mats_tb_intrinsic)
mats_extrinsic = (mats_euler_extrinsic, mats_tb_extrinsic)
# The beta angles are different for proper Euler angles and Tait-Bryan
# angles because the relevant beta terms will be sines and cosines,
# respectively.
all_betas = ((0, np.pi), (np.pi / 2, -np.pi / 2))
alphas = (0, np.pi / 2, np.pi, 3 * np.pi / 2)
for mats in (mats_intrinsic, mats_extrinsic):
for betas, mat_set in zip(all_betas, mats):
for convention, axis_type, mat_func in mat_set:
quaternions = []
for beta in betas:
for alpha in alphas:
mat = mat_func(alpha, beta)
if np.linalg.det(mat) == -1:
# Some of these will be improper rotations.
continue
quat = rowan.from_matrix(mat)
quaternions.append(quat)
euler = rowan.to_euler(quat, convention, axis_type)
converted = rowan.from_euler(*euler,
convention=convention,
axis_type=axis_type)
correct_rotation = rowan.rotate(quat, test_vector)
test_rotation = rowan.rotate(
converted, test_vector)
self.assertTrue(np.allclose(correct_rotation,
test_rotation,
atol=1e-6),
msg="""
Failed for convention {},
axis type {},
alpha = {},
beta = {}.
Expected quaternion: {}.
Calculated: {}.
Expected vector: {}.
Calculated vector: {}.""".format(
convention, axis_type, alpha,
beta, quat, converted,
correct_rotation,
test_rotation))
# For completeness, also test with broadcasting.
quaternions = np.asarray(quaternions).reshape(-1, 4)
all_euler = rowan.to_euler(quaternions, convention,
axis_type)
converted = rowan.from_euler(all_euler[...,
0], all_euler[...,
1],
all_euler[..., 2], convention,
axis_type)
self.assertTrue(
np.allclose(rowan.rotate(quaternions, test_vector),
rowan.rotate(converted, test_vector),
atol=1e-6))
def policy(self,
X,
pd=np.zeros(3),
vd=np.zeros(3),
ad=np.zeros(3),
jd=np.zeros(3),
sd=np.zeros(3),
logentry=None,
time=None,
K_i=0,
K_p=0,
K_d=0):
# print('baseline: X = ', X)
p_e = pd - X[:3]
v_e = vd - X[3:6]
print(X)
int_p_e = self.integrate_error(p_e, time)
F_d = (self.K_i * int_p_e + self.K_p * p_e +
self.K_d * v_e) * self.mass # TODO: add integral term
F_d[2] += self.g * self.mass
T_d = np.linalg.norm(F_d)
yaw_d = 0
roll_d = np.arcsin(
(F_d[0] * np.sin(yaw_d) - F_d[1] * np.cos(yaw_d)) / T_d)
pitch_d = np.arctan(
(F_d[0] * np.cos(yaw_d) + F_d[1] * np.sin(yaw_d)) / F_d[2])
euler_d = np.array([roll_d, pitch_d, yaw_d])
q = rowan.normalize(X[6:10])
euler = rowan.to_euler(q, 'xyz')
att_e = -(euler - euler_d)
th_r = rowan.normalize(
rowan.from_euler(att_e[0], att_e[1], att_e[2], 'xyz'))
th_d = rowan.normalize(
rowan.from_euler(euler_d[0], euler_d[1], euler_d[2], 'xyz'))
att_v_e = -X[10:]
torque = self.Att_p * att_e + self.Att_d * att_v_e
torque[0] *= self.J[0]
torque[1] *= self.J[1]
torque[2] *= self.J[2]
Jomega = np.array(
[self.J[0] * X[10], self.J[1] * X[11], self.J[2] * X[12]])
torque -= np.cross(Jomega, X[10:])
yawpart = -0.25 * torque[2] / self.t2t
rollpart = 0.25 / self.arm * torque[0] #/ self.t2t
pitchpart = 0.25 / self.arm * torque[1] #/ self.t2t
thrustpart = 0.25 * T_d
motorForce = np.array([
thrustpart - rollpart - pitchpart + yawpart,
thrustpart - rollpart + pitchpart - yawpart,
thrustpart + rollpart + pitchpart + yawpart,
thrustpart + rollpart - pitchpart - yawpart
])
Fc = motorForce / self.params['C_T']
omega = np.sqrt(np.maximum(Fc, self.params['motor_min_speed'])
) # Ensure each prop has a positive thrust
omega = np.minimum(
omega, self.params['motor_max_speed']) # Maximum rotation speed
# print('baseline: omega = ', omega)
# print('T', T_r)
# print('tau',tau_r)
# x_error = pd[0] - X[0]
# y_error = pd[1] - X[1]
# z_error = pd[2] - X[2]
#
# vx_error = vd[0] - X[3]
# vy_error = vd[1] - X[4]
# vz_error = vd[2] - X[5]
#
# # print('baseline: x_error = ', x_error)
# # print('baseline: y_error = ', y_error)
#
# ax_d = ad[0]
# ay_d = ad[1]
# az_d = ad[2]
#
# self.int_x_error += self.params.dt_posctrl * x_error
# self.int_y_error += self.params.dt_posctrl * y_error
# self.int_z_error += self.params.dt_posctrl * z_error
#
# ax_r = self.params['K_p_x'] * x_error + self.params['K_d_x'] * vx_error + self.params[
# 'K_i_x'] * self.int_x_error + ax_d
# ay_r = self.params['K_p_y'] * y_error + self.params['K_d_y'] * vy_error + self.params[
# 'K_i_y'] * self.int_y_error + ay_d
# az_r = self.params['K_p_z'] * z_error + self.params['K_d_z'] * vz_error + self.params[
# 'K_i_z'] * self.int_z_error + az_d
# # print('baseline: x PDI terms:',K_p_x * x_error, K_d_x * vx_error, K_i_x * intx_error, ax_d)
# # print('baseline: y PDI terms:',K_p_y * y_error, K_d_y * vy_error, K_i_y * inty_error, ay_d)
#
# jx_d = jd[0]
# jy_d = jd[1]
# jz_d = jd[2]
# # jx_r = K_p_x * vx_error + K_d_x * (ax_d - ax_r) + K_i_x * x_error + jx_d
# # jy_r = K_p_y * vy_error + K_d_y * (ay_d - ay_r) + K_i_y * y_error + jy_d
#
# sx_d = sd[0]
# sy_d = sd[1]
# sz_d = sd[2]
# # sx_r = K_p_x * (ax_d - ax_r) + K_d_x * (jx_d - jx_r) + K_i_x * vx_error + sx_d
# # sy_r = K_p_y * (ay_d - ay_r) + K_d_y * (jy_d - jy_r) + K_i_y * vy_error + sy_d
#
# # print('baseline: ay_r = ', ay_r)
#
# Fx_r = ax_r * self.params['m']
# Fy_r = (ay_r + self.params['g']) * self.params['m']
# Fz_r = az_r * self.params['m']
#
# Fy_r = np.maximum(Fy_r, 0.10 * self.params['g'] * self.params['m'])
#
#
# # print('baseline: Fx_r, Fy_r = ', Fx_r, Fy_r)
#
# T_r = np.minimum(np.sqrt(Fy_r ** 2 + Fx_r ** 2 + Fz_r ** 2), 125.) # / np.cos(X[2])
# if T_r > 124.:
# warn('thrust gets too high')
# # print('T_r = ', T_r)
# T_d = self.params['m'] * np.sqrt(ax_d ** 2 + (ay_d + self.params['g']) ** 2 + az_d ** 2)
#
# # th = X[2]
# th_r =
# th_d =
#
# mat = np.array([[-np.sin(th_d), -T_d * np.cos(th_d)], #TODO: use th_d, T_d
# [np.cos(th_d), -T_d * np.sin(th_d)]])
# b = self.params['m'] * np.array([jx_d, jy_d, jz_d])#
# v = np.linalg.solve(mat, b)
# T_d_dot = v[0]
# w_d = v[1]
# # w = X[5]
#
# mat = np.array([[-np.sin(th_d), -T_d * np.cos(th_d)],
# [np.cos(th_d), -T_d * np.sin(th_d)]])
# b = self.params['m'] * np.array([sx_d, sy_d]) - np.array(
# [-2. * T_d_dot * np.cos(th_d) * w_d + T_d * np.sin(th_d) * (w_d ** 2),
# -2. * T_d_dot * np.sin(th_d) * w_d - T_d * np.cos(th_d) * (w_d ** 2)])
#
# v = np.linalg.solve(mat, b)
# # T_d_ddot = v[0]
# alpha_d = v[1]
# # print('baseline_pos: Tdot, w_d, alpha_d', Tdot, w_d, alpha_d)
w_d = 0
alpha_d = 0
if logentry is not None:
logentry['th_d'] = th_d
logentry['th_r'] = th_r
logentry['T_d'] = T_d
logentry['euler'] = euler
logentry['att_e'] = att_e
# logentry['Fx_r'] = Fx_r
# logentry['Fy_r'] = Fy_r
# logentry['Fz_r'] = Fz_r
# logentry['int_x_error'] = self.int_x_error
# logentry['int_y_error'] = self.int_y_error
# logentry['int_y_error'] = self.int_z_error
return T_d, th_d, w_d, alpha_d, omega, motorForce
def policy(self, state): # current state p = state[0:3] v = state[3:6] q = state[6:10] omega = state[10:13] R = rowan.to_matrix(q) # position controller # p_tilde = p - self.p_d # v_tilde = v - self.v_d # v_r = self.v_d - self.lambda_p * p_tilde # v_r_dot = self.a_d - self.lambda_p * v_tilde # s_v = v - v_r # f_r = self.mass * v_r_dot \ # - self.Kv * s_v \ # - self.Kp * p_tilde # thrust = np.linalg.norm(f_r) # qr = mkvec( # asinf((F_d.x * sinf(yaw) - F_d.y * cosf(yaw)) / control->thrustSI), # atanf((F_d.x * cosf(yaw) + F_d.y * sinf(yaw)) / F_d.z), # desiredYaw); # position controller pos_e = self.p_d - p vel_e = self.v_d - v F_d = self.mass * (self.a_d + self.Kpos_D * vel_e + self.Kpos_P * pos_e) print(F_d) thrust = np.linalg.norm(F_d) yaw = 0 rpy_d = np.array([ np.arcsin((F_d[0] * np.sin(yaw) - F_d[1] * np.cos(yaw)) / thrust), np.arctan((F_d[0] * np.cos(yaw) + F_d[1] * np.sin(yaw)) / F_d[2]), yaw]) R_d = rowan.to_matrix(rowan.from_euler(*rpy_d)) print(rpy_d) print(R_d) # attitude controller # rotation error Rtilde = self.R_d.T @ R qtilde_0 = 1/2 * np.sqrt(1 + np.trace(Rtilde)) qtilde_v = 1 / (4 * qtilde_0) * veemap(Rtilde - Rtilde.T) omega_r = Rtilde.T @ self.omega_d - 2 * self.lambda_a * qtilde_v print(omega_r) if self.omega_r_last is not None: omega_r_dot = (omega_r - self.omega_r_last) / 0.01 else: omega_r_dot = np.zeros(3) self.omega_r_last = omega_r s_omega = omega - omega_r torque = self.J * omega_r_dot \ - np.cross(self.J * omega, omega_r) \ - self.K_att * s_omega \ - self.kq * qtilde_v # power distribution thrust_to_torque = 0.006 arm_length = 0.046 thrustpart = 0.25 * thrust yawpart = -0.25 * torque[2] / thrust_to_torque arm = 0.707106781 * arm_length rollpart = 0.25 / arm * torque[0] pitchpart = 0.25 / arm * torque[1] motorForce = np.array([ thrustpart - rollpart - pitchpart + yawpart, thrustpart - rollpart + pitchpart - yawpart, thrustpart + rollpart + pitchpart + yawpart, thrustpart + rollpart - pitchpart - yawpart ]) motorForce = np.clip(motorForce, self.a_min, self.a_max) # logging self.qr.append(np.degrees(rpy_d)) self.q.append(np.degrees(rowan.to_euler(q, 'xyz'))) self.omegar.append(omega_r) self.omega.append(omega) return motorForce
def __init__(self):
super().__init__()
self.env_name = 'Quadrotor'
self.env_case = 'SmallAngle'
# flags
self.rl_continuous_on = True
self.sim_render_on = False
self.pomdp_on = False
self.single_agent_sim = True
self.multi_agent_sim = False
# Crazyflie 2.0 quadrotor
self.mass = 0.034 # kg
# self.J = np.array([
# [16.56,0.83,0.71],
# [0.83,16.66,1.8],
# [0.72,1.8,29.26]
# ]) * 1e-6 # kg m^2
self.J = np.array([16.571710e-6, 16.655602e-6, 29.261652e-6])
# Note: we assume here that our control is forces
arm_length = 0.046 # m
arm = 0.707106781 * arm_length
t2t = 0.006 # thrust-to-torque ratio
self.B0 = np.array([[1, 1, 1, 1], [-arm, -arm, arm, arm],
[-arm, arm, arm, -arm], [-t2t, t2t, -t2t, t2t]])
self.g = 9.81 # not signed
# control limits [N]
self.a_min = np.array([0, 0, 0, 0])
self.a_max = np.array([12, 12, 12, 12]) / 1000 * 9.81 # g->N
# perfect hover would use: np.array([0.0085, 0.0085, 0.0085, 0.0085]) * 9.81
# self.a_min = np.array([0.008, 0.008, 0.008, 0.008]) * 9.81
# self.a_max = np.array([0.012, 0.012, 0.012, 0.012]) * 9.81 # g->N
# RL
self.rl_train_model_fn = '../models/quadrotor/rl_current.pt'
self.rl_lr_schedule_on = False
self.rl_lr_schedule_gamma = 0.2
self.rl_warm_start_on = False
self.rl_warm_start_fn = '../models/quadrotor/rl_continuous_v3.pt'
self.rl_module = 'DDPG'
self.rl_lr_schedule = np.arange(0, 10)
# common param
self.rl_gamma = 0.999
self.rl_K = 10
self.rl_max_episodes = 50000
self.rl_batch_size = 2000
if self.rl_continuous_on:
# ddpg param
self.rl_lr_mu = 1e-4
self.rl_lr_q = 1e-3
self.rl_buffer_limit = 5e6
self.rl_action_std = 0.05
self.rl_max_action_perturb = 0.05
self.rl_tau = 0.995
# network architecture
n, m, h_mu, h_q = 13, 4, 64, 64 # state dim, action dim, hidden layers
self.rl_mu_network_architecture = nn.ModuleList([
nn.Linear(n, h_mu),
nn.Linear(h_mu, h_mu),
nn.Linear(h_mu, m)
])
self.rl_q_network_architecture = nn.ModuleList([
nn.Linear(n + m, h_q),
nn.Linear(h_q, h_q),
nn.Linear(h_q, 1)
])
self.rl_network_activation = tanh
else:
# ppo param s
self.rl_lr = 5e-3
self.rl_lmbda = 0.95
self.rl_eps_clip = 0.2
self.rl_discrete_action_space = [
np.array([0, 0, 0, 0]) * 12 / 1000 * 9.81,
np.array([0, 0, 0, 1]) * 12 / 1000 * 9.81,
np.array([0, 0, 1, 0]) * 12 / 1000 * 9.81,
np.array([0, 0, 1, 1]) * 12 / 1000 * 9.81,
np.array([0, 1, 0, 0]) * 12 / 1000 * 9.81,
np.array([0, 1, 0, 1]) * 12 / 1000 * 9.81,
np.array([0, 1, 1, 0]) * 12 / 1000 * 9.81,
np.array([0, 1, 1, 1]) * 12 / 1000 * 9.81,
np.array([1, 0, 0, 0]) * 12 / 1000 * 9.81,
np.array([1, 0, 0, 1]) * 12 / 1000 * 9.81,
np.array([1, 0, 1, 0]) * 12 / 1000 * 9.81,
np.array([1, 0, 1, 1]) * 12 / 1000 * 9.81,
np.array([1, 1, 0, 0]) * 12 / 1000 * 9.81,
np.array([1, 1, 0, 1]) * 12 / 1000 * 9.81,
np.array([1, 1, 1, 0]) * 12 / 1000 * 9.81,
np.array([1, 1, 1, 1]) * 12 / 1000 * 9.81,
]
self.rl_lr = 1e-3 #5e-3
# IL
self.il_train_model_fn = '../models/quadrotor/il_current.pt'
self.il_imitate_model_fn = '../models/quadrotor/rl_current.pt'
# Sim
self.sim_rl_model_fn = '../models/quadrotor/rl_current.pt' # rl_current
self.sim_il_model_fn = '../models/quadrotor/il_current.pt'
self.sim_t0 = 0
self.sim_tf = 3
self.sim_dt = 0.01
self.sim_times = np.arange(self.sim_t0, self.sim_tf, self.sim_dt)
self.sim_nt = len(self.sim_times)
s_desired = np.zeros(13)
s_desired[6:10] = rowan.from_euler(np.radians(0), np.radians(0),
np.radians(0), 'xyz')
self.ref_trajectory = np.tile(
np.array([s_desired.T]).T, (1, self.sim_nt))
def rotation_z(num: int, angle):
return rowan.from_euler(np.ones(num) * angle, np.zeros(num), np.zeros(num))
