def polar_axis_angle_to_quaternion(polar_axis_angle):
"""
Convert polar axis-angle representation, which constitutes the elevation and azimuth angles of the axis as well
as the rotation angle :math:`\mathbf{θ}_{paa} = [ϕ_e, ϕ_a, θ]`, to quaternion form
:math:`\mathbf{q} = [q_i, q_j, q_k, q_r]`.\n
`[reference] <https://en.wikipedia.org/wiki/Axis%E2%80%93angle_representation>`_
:param polar_axis_angle: Polar axis angle representation *[batch_shape,3]*
:type polar_axis_angle: array
:return: Quaternion *[batch_shape,4]*
"""
# BS x 1
theta = polar_axis_angle[..., 0:1]
phi = polar_axis_angle[..., 1:2]
angle = polar_axis_angle[..., 2:3]
x = _ivy.sin(theta) * _ivy.cos(phi)
y = _ivy.sin(theta) * _ivy.sin(phi)
z = _ivy.cos(theta)
# BS x 3
vector = _ivy.concatenate((x, y, z), -1)
# BS x 4
return axis_angle_to_quaternion(_ivy.concatenate([vector, angle], -1))
def get_reward(self):
"""
Get reward based on current state
:return: Reward array
"""
# Pole verticality.
rew = (ivy.cos(self.angle) + 1) / 2
return ivy.reshape(rew, (1, ))
def get_observation(self):
"""
Get observation from environment.
:return: observation array
"""
return ivy.concatenate(
(ivy.cos(self.angle), ivy.sin(self.angle), self.angle_vel),
axis=-1)
def get_reward(self):
"""
Get reward based on current state
:return: Reward array
"""
# Center proximity.
rew = ivy.exp(-1 * (self.x**2))
# Pole verticality.
rew = rew * (ivy.cos(self.angle) + 1) / 2
return ivy.reshape(rew[0], (1, ))
def test_cos(x, dtype_str, tensor_fn, dev_str, call):
# smoke test
x = tensor_fn(x, dtype_str, dev_str)
ret = ivy.cos(x)
# type test
assert ivy.is_array(ret)
# cardinality test
assert ret.shape == x.shape
# value test
assert np.allclose(call(ivy.cos, x), ivy.numpy.cos(ivy.to_numpy(x)))
# compilation test
helpers.assert_compilable(ivy.cos)
def get_observation(self):
"""
Get observation from environment.
:return: observation array
"""
ob = (ivy.reshape(ivy.cos(self.angles),
(1, 2)), ivy.reshape(ivy.sin(self.angles), (1, 2)),
ivy.reshape(self.angle_vels,
(1, 2)), ivy.reshape(self.goal_xy, (1, 2)))
ob = ivy.concatenate(ob, axis=0)
return ivy.reshape(ob, (-1, ))
def polar_to_cartesian_coords(polar_coords):
"""
Convert spherical polar co-ordinates :math:`\mathbf{x}_p = [r, α, β]` to cartesian co-ordinates
:math:`\mathbf{x}_c = [x, y, z]`.\n
`[reference] <https://en.wikipedia.org/wiki/Spherical_coordinate_system#Cartesian_coordinates>`_
:param polar_coords: Spherical polar co-ordinates *[batch_shape,3]*
:type polar_coords: array
:return: Cartesian co-ordinates *[batch_shape,3]*
"""
# BS x 1
phi = polar_coords[..., 0:1]
theta = polar_coords[..., 1:2]
r = polar_coords[..., 2:3]
x = r * _ivy.sin(theta) * _ivy.cos(phi)
y = r * _ivy.sin(theta) * _ivy.sin(phi)
z = r * _ivy.cos(theta)
# BS x 3
return _ivy.concatenate((x, y, z), -1)
def get_reward(self):
"""
Get reward based on current state
:return: Reward array
"""
# Goal proximity.
x = ivy.reduce_sum(ivy.cos(self.angles), -1)
y = ivy.reduce_sum(ivy.sin(self.angles), -1)
xy = ivy.concatenate([ivy.expand_dims(x, 0),
ivy.expand_dims(y, 0)],
axis=0)
rew = ivy.reshape(
ivy.exp(-1 * ivy.reduce_sum((xy - self.goal_xy)**2, -1)), (-1, ))
return ivy.reduce_mean(rew, axis=0, keepdims=True)
def step(self, action):
force = self.torque_scale * action
angle_cos = ivy.cos(self.angle)
angle_sin = ivy.sin(self.angle)
temp = (
(force + self.pole_mass_length * self.angle_vel**2 * angle_sin) /
self.total_mass)
angle_acc = (
(self.g * angle_sin - angle_cos * temp) /
(self.pole_length *
(4.0 / 3.0 - self.pole_mass * angle_cos**2 / self.total_mass)))
x_acc = (
temp -
self.pole_mass_length * angle_acc * angle_cos / self.total_mass)
self.x_vel = self.x_vel + self.dt * x_acc
self.x = self.x + self.dt * self.x_vel
self.angle_vel = self.angle_vel + self.dt * angle_acc
self.angle = self.angle + self.dt * self.angle_vel
return self.get_observation(), self.get_reward(), False, {}
def axis_angle_to_quaternion(axis_angle):
"""
Convert rotation axis unit vector :math:`\mathbf{e} = [e_x, e_y, e_z]` and rotation angle :math:`θ` to quaternion
:math:`\mathbf{q} = [q_i, q_j, q_k, q_r]`.\n
`[reference] <https://en.wikipedia.org/wiki/Rotation_formalisms_in_three_dimensions#Quaternions>`_
:param axis_angle: Axis to rotate about and angle to rotate *[batch_shape,4]*
:type axis_angle: array
:return: Quaternion *[batch_shape,4]*
"""
# BS x 1
angle = axis_angle[..., -1:]
n = _ivy.cos(angle / 2)
e1 = _ivy.sin(angle / 2) * axis_angle[..., 0:1]
e2 = _ivy.sin(angle / 2) * axis_angle[..., 1:2]
e3 = _ivy.sin(angle / 2) * axis_angle[..., 2:3]
# BS x 4
quaternion = _ivy.concatenate((e1, e2, e3, n), -1)
return quaternion
def render(self, mode='human'):
"""
Renders the environment.
The set of supported modes varies per environment. (And some
environments do not support rendering at all.) By convention,
if mode is:
- human: render to the current display or terminal and
return nothing. Usually for human consumption.
- rgb_array: Return an numpy.ndarray with shape (x, y, 3),
representing RGB values for an x-by-y pixel image, suitable
for turning into a video.
- ansi: Return a string (str) or StringIO.StringIO containing a
terminal-style text representation. The text can include newlines
and ANSI escape sequences (e.g. for colors).
:param mode: Render mode, one of [human|rgb_array], default human
:type mode: str, optional
:return: Rendered image.
"""
if self.viewer is None:
# noinspection PyBroadException
try:
from gym.envs.classic_control import rendering
except:
if not self._logged_headless_message:
print(
'Unable to connect to display. Running the Ivy environment in headless mode...'
)
self._logged_headless_message = True
return
self.viewer = rendering.Viewer(500, 500)
bound = self.num_joints + 0.2
self.viewer.set_bounds(-bound, bound, -bound, bound)
# Goal.
goal_geom = rendering.make_circle(0.2)
goal_geom.set_color(0.4, 0.6, 1.)
self.goal_tr = rendering.Transform()
goal_geom.add_attr(self.goal_tr)
self.viewer.add_geom(goal_geom)
# Arm segments and joints.
l, r, t, b = 0, 1., 0.1, -0.1
self.segment_trs = []
for _ in range(self.num_joints):
# Segment.
segment_geom = rendering.FilledPolygon([(l, b), (l, t), (r, t),
(r, b)])
segment_tr = rendering.Transform()
self.segment_trs.append(segment_tr)
segment_geom.add_attr(segment_tr)
segment_geom.set_color(0., 0., 0.)
self.viewer.add_geom(segment_geom)
# Joint.
joint_geom = rendering.make_circle(0.1)
joint_geom.set_color(0.5, 0.5, 0.5)
joint_geom.add_attr(segment_tr)
self.viewer.add_geom(joint_geom)
# End effector.
self.end_geom = rendering.make_circle(0.1)
self.end_tr = rendering.Transform()
self.end_geom.add_attr(self.end_tr)
self.viewer.add_geom(self.end_geom)
self.goal_tr.set_translation(*ivy.to_numpy(self.goal_xy).tolist())
x, y = 0., 0.
for segment_tr, angle in zip(self.segment_trs,
ivy.reshape(self.angles, (-1, 1))):
segment_tr.set_rotation(ivy.to_numpy(angle)[0])
segment_tr.set_translation(x, y)
x = ivy.to_numpy(x + ivy.cos(ivy.expand_dims(angle, 0))[0])[0]
y = ivy.to_numpy(y + ivy.sin(ivy.expand_dims(angle, 0))[0])[0]
self.end_tr.set_translation(x, y)
rew = ivy.to_numpy(self.get_reward())[0]
self.end_geom.set_color(1 - rew, rew, 0.)
return self.viewer.render(return_rgb_array=mode == 'rgb_array')
def render(self, mode='human'):
"""
Renders the environment.
The set of supported modes varies per environment. (And some
environments do not support rendering at all.) By convention,
if mode is:
- human: render to the current display or terminal and
return nothing. Usually for human consumption.
- rgb_array: Return an numpy.ndarray with shape (x, y, 3),
representing RGB values for an x-by-y pixel image, suitable
for turning into a video.
- ansi: Return a string (str) or StringIO.StringIO containing a
terminal-style text representation. The text can include newlines
and ANSI escape sequences (e.g. for colors).
:param mode: Render mode, one of [human|rgb_array], default human
:type mode: str, optional
:return: Rendered image.
"""
screen_width = 500
screen_height = 500
x_min = -1.2
x_max = 0.6
world_width = x_max - x_min
scale = screen_width / world_width
car_width = 40
car_height = 20
if self.viewer is None:
# noinspection PyBroadException
try:
from gym.envs.classic_control import rendering
except:
if not self._logged_headless_message:
print('Unable to connect to display. Running the Ivy environment in headless mode...')
self._logged_headless_message = True
return
self.viewer = rendering.Viewer(screen_width, screen_height)
# Track.
xs = ivy.linspace(x_min, x_max, 100)
ys = self._height(xs)
xys = list((ivy.to_numpy(xt).item(), ivy.to_numpy(yt).item())
for xt, yt in zip((xs - x_min) * scale, ys * scale))
self.track = rendering.make_polyline(xys)
self.track.set_linewidth(2)
self.viewer.add_geom(self.track)
# Car.
clearance = 10
l, r, t, b = -car_width / 2, car_width / 2, car_height, 0
self.car_geom = rendering.FilledPolygon(
[(l, b), (l, t), (r, t), (r, b)])
self.car_geom.add_attr(
rendering.Transform(translation=(0, clearance)))
self.car_tr = rendering.Transform()
self.car_geom.add_attr(self.car_tr)
self.viewer.add_geom(self.car_geom)
# Wheels.
front_wheel = rendering.make_circle(car_height / 2.5)
front_wheel.set_color(0.5, 0.5, 0.5)
front_wheel.add_attr(
rendering.Transform(translation=(car_width / 4, clearance)))
front_wheel.add_attr(self.car_tr)
self.viewer.add_geom(front_wheel)
back_wheel = rendering.make_circle(car_height / 2.5)
back_wheel.add_attr(
rendering.Transform(translation=(-car_width / 4, clearance)))
back_wheel.add_attr(self.car_tr)
back_wheel.set_color(0.5, 0.5, 0.5)
self.viewer.add_geom(back_wheel)
# Flag.
flag_x = (ivy.to_numpy(self.goal_x)[0] - x_min) * scale
flagy_y1 = ivy.to_numpy(self._height(self.goal_x))[0] * scale
flagy_y2 = flagy_y1 + 50
flagpole = rendering.Line((flag_x, flagy_y1), (flag_x, flagy_y2))
self.viewer.add_geom(flagpole)
flag = rendering.FilledPolygon(
[(flag_x, flagy_y2), (flag_x, flagy_y2 - 10),
(flag_x + 25, flagy_y2 - 5)])
flag.set_color(0.4, 0.6, 1.)
self.viewer.add_geom(flag)
self.car_tr.set_translation(
(ivy.to_numpy(self.x)[0] - x_min) * scale, ivy.to_numpy(self._height(self.x))[0] * scale)
self.car_tr.set_rotation(ivy.to_numpy(ivy.cos(3 * self.x))[0])
rew = ivy.to_numpy(self.get_reward()).item()
self.car_geom.set_color(1 - rew, rew, 0.)
return self.viewer.render(return_rgb_array=mode == 'rgb_array')
def step(self, action):
x_acc = action * self.torque_scale - self.g * ivy.cos(3 * self.x)
self.x_vel = self.x_vel + self.dt * x_acc
self.x = self.x + self.dt * self.x_vel
return self.get_observation(), self.get_reward(), False, {}
def __init__(self, a_s, d_s, alpha_s, dh_joint_scales, dh_joint_offsets, base_inv_ext_mat=None):
"""
Initialize robot manipulator instance
:param a_s: Denavit–Hartenberg "a" parameters *[num_joints]*
:type a_s: array
:param d_s: Denavit–Hartenberg "d" parameters *[num_joints]*
:type d_s: array
:param alpha_s: Denavit–Hartenberg "alpha" parameters *[num_joints]*
:type alpha_s: array
:param dh_joint_scales: Scalars to apply to input joints *[num_joints]*
:type dh_joint_scales: array
:param dh_joint_offsets: Scalar offsets to apply to input joints *[num_joints]*
:type dh_joint_offsets: array
:param base_inv_ext_mat: Inverse extrinsic matrix of the robot base *[4,4]*
:type base_inv_ext_mat: array, optional
"""
self._num_joints = a_s.shape[-1]
# ToDo: incorporate the base_inv_ext_mat more elegantly, instead of the hack as in the sample_links method
if base_inv_ext_mat is None:
self._base_inv_ext_mat = _ivy.identity(4)
else:
self._base_inv_ext_mat = base_inv_ext_mat
# NJ
self._dis = d_s
self._dh_joint_scales = dh_joint_scales
self._dh_joint_offsets = dh_joint_offsets
# Forward Kinematics Constant Matrices
# Based on Denavit-Hartenberg Convention
# Using same nomenclature as in:
# Modelling, Planning and Control. Bruno Siciliano, Lorenzo Sciavicco, Luigi Villani, Giuseppe Oriolo
# page 61 - 65
AidashtoAis_list = [_ivy.identity(4, batch_shape=[1])]
# repeated blocks
# 1 x 1 x 3
start_of_top_row = _ivy.array([[[1., 0., 0.]]])
# 1 x 1 x 1
zeros = _ivy.zeros((1, 1, 1))
# 1 x 1 x 4
bottom_row = _ivy.array([[[0., 0., 0., 1.]]])
for i in range(self._num_joints):
# 1 x 1 x 1
a_i = _ivy.reshape(a_s[i], [1] * 3)
cos_alpha = _ivy.reshape(_ivy.cos(alpha_s[i]), [1] * 3)
sin_alpha = _ivy.reshape(_ivy.sin(alpha_s[i]), [1] * 3)
# 1 x 1 x 4
top_row = _ivy.concatenate((start_of_top_row, a_i), -1)
top_middle_row = _ivy.concatenate((zeros, cos_alpha, -sin_alpha, zeros), -1)
bottom_middle_row = _ivy.concatenate((zeros, sin_alpha, cos_alpha, zeros), -1)
# 1 x 4 x 4
AidashtoAi = _ivy.concatenate((top_row, top_middle_row, bottom_middle_row, bottom_row), 1)
# list
AidashtoAis_list.append(AidashtoAi)
# NJ x 4 x 4
self._AidashtoAis = _ivy.concatenate(AidashtoAis_list, 0)
# Constant Jacobian Params
# Using same nomenclature as in:
# Modelling, Planning and Control. Bruno Siciliano, Lorenzo Sciavicco, Luigi Villani, Giuseppe Oriolo
# page 111 - 113
# 1 x 3
self._z0 = _ivy.array([[0],
[0],
[1]])
# 1 x 4
self._p0hat = _ivy.array([[0],
[0],
[0],
[1]])
# link lengths
# NJ
self._link_lengths = (a_s ** 2 + d_s ** 2) ** 0.5
def compute_link_matrices(self, joint_angles, link_num, batch_shape=None):
"""
Compute homogeneous transformation matrices relative to base frame, up to link_num of links.
:param joint_angles: Joint angles of the robot *[batch_shape,num_joints]*
:type joint_angles: array
:param link_num: Link number for which to compute matrices up to
:type link_num: int
:param batch_shape: Shape of batch. Inferred from inputs if None.
:type batch_shape: sequence of ints, optional
:return: The link_num matrices, up the link_num *[batch_shape,link_num,4,4]*
"""
if batch_shape is None:
batch_shape = joint_angles.shape[:-1]
batch_shape = list(batch_shape)
num_batch_dims = len(batch_shape)
# BS x 1 x NJ
try:
dh_joint_angles = _ivy.expand_dims(joint_angles * self._dh_joint_scales - self._dh_joint_offsets, -2)
except:
d = 0
# BS x 1 x 4 x 4
A00 = _ivy.identity(4, batch_shape=batch_shape + [1])
Aitoip1dashs = list()
Aiip1s = list()
A0is = [A00]
# repeated blocks
# BS x 1 x NJ
dis = _ivy.tile(_ivy.reshape(self._dis, [1] * num_batch_dims + [1, self._num_joints]),
batch_shape + [1, 1])
# BS x 1 x 4
bottom_row = _ivy.tile(
_ivy.reshape(_ivy.array([0., 0., 0., 1.]), [1] * num_batch_dims + [1, 4]),
batch_shape + [1, 1])
# BS x 1 x 3
start_of_bottom_middle = _ivy.tile(
_ivy.reshape(_ivy.array([0., 0., 1.]), [1] * num_batch_dims + [1, 3]),
batch_shape + [1, 1])
# BS x 1 x 2
zeros = _ivy.zeros(batch_shape + [1, 2])
for i in range(self._num_joints):
# BS x 1 x 4
top_row = _ivy.concatenate((_ivy.cos(dh_joint_angles[..., i:i + 1]),
-_ivy.sin(dh_joint_angles[..., i:i + 1]), zeros), -1)
top_middle_row = _ivy.concatenate((_ivy.sin(dh_joint_angles[..., i:i + 1]),
_ivy.cos(dh_joint_angles[..., i:i + 1]), zeros), -1)
bottom_middle_row = _ivy.concatenate((start_of_bottom_middle, dis[..., i:i + 1]), -1)
# BS x 4 x 4
Aitoip1dash = _ivy.concatenate((top_row, top_middle_row, bottom_middle_row, bottom_row), -2)
# (BSx4) x 4
Aitoip1dash_flat = _ivy.reshape(Aitoip1dash, (-1, 4))
# (BSx4) x 4
Aiip1_flat = _ivy.matmul(Aitoip1dash_flat, self._AidashtoAis[i + 1])
# BS x 4 x 4
Aiip1 = _ivy.reshape(Aiip1_flat, batch_shape + [4, 4])
# BS x 4 x 4
A0ip1 = _ivy.matmul(A0is[-1][..., 0, :, :], Aiip1)
# append term to lists
Aitoip1dashs.append(Aitoip1dash)
Aiip1s.append(Aiip1)
A0is.append(_ivy.expand_dims(A0ip1, -3))
if i + 1 == link_num:
# BS x LN x 4 x 4
return _ivy.concatenate(A0is, -3)
raise Exception('wrong parameter entered for link_num, please enter integer from 1-' + str(self._num_joints))
