python - TensorFlow 中的变换矩阵

Question

我想用 TensorFlow 实现机器人的正向运动学；主要是为了获得自动微分并将该模块插入到更大的网络架构中。

一般来说，我有一堆 4x4 变换矩阵，由 dh 参数定义(d、theta、a、>alpha)和关节角度q:

[[           cos(theta+q),           -sin(theta+q),           0,             a],
 [sin(theta+q)*cos(alpha), cos(theta+q)*cos(alpha), -sin(alpha), -sin(alpha)*d],
 [sin(theta+q)*sin(alpha), cos(theta+q)*sin(alpha),  cos(alpha),  cos(alpha)*d],
 [                      0,                       0,           0,             1]])

我的机器人有 10 个不同的关节，全部按顺序连接。 我认为预先计算正弦和余弦会很聪明。

q = tf.keras.layers.Input((10,))
sin_q = tf.sin(q)
cos_q = tf.cos(q)

让我们看看第一个关节处具有特定 dh 参数集的转换(d=0.1055、theta=0、a=0，alpha=0):

m0 = [[cos(q0), -sin(q0), 0,       0],
      [sin(q0),  cos(q0), 0,       0],
            0,       0,   1, 0.10550],
            0,       0,   0,       1]] 

我的第一个问题是如何使用 TensorFlow 构建这样的东西？ 在 numpy 中，我将初始化矩阵并填充非零值。

m_shape = tf.TensorShape((batch_size,4,4))
m0 = tf.zeros(m_shape)
m0[..., 0, 0] = cos_q[..., 0]
m0[..., 0, 1] = -sin_q[..., 0]
m0[..., 1, 0] = cos_q[..., 0]
m0[..., 1, 1] = sin_q[..., 0]
m0[..., 2, 3] = 0.10550
m0[..., 3, 3] = 1
Error -> 'Tensor' object does not support item assignment

但是 Tensorflow 不允许对张量进行赋值。 看来要走的路是通过 tf.stack() 。我需要创建一个与我未指定的batch_size、堆栈和 reshape 大小相同的向量。 (注:一般情况下零值较少)

e = tf.ones_like(q[..., 0]) 
m0 = tf.stack([cos_q[..., 0], -sin_q[..., 0], 0*e,       0*e,
               sin_q[..., 0],  cos_q[..., 0], 0*e,       0*e,
                         0*e,            0*e, 1*e, 0.10550*e,
                         0*e,            0*e, 0*e,       1*e], axis=-1)
m0 = tf.keras.layers.Reshape((4, 4))(m0)

这是正确的还是有更智能的方法在 TensorFlow 中构建此类通用转换？

作为最终结果，我对运动链末端的转换感兴趣。我想放入一组不同的关节配置 (?, 10) 并在末端执行器 (?, 4, 4) 处进行转换。

m_end = m0 @ m1 @ m2 @ ... @ m10
forward_net = tf.keras.Model(inputs=[q], outputs=[m_end]
result = forward_net.predict(np.random.random((100, 10)))

这可行，但既不优雅也不快速。 速度是我最大的问题； numpy 中的相同实现速度快了 150 倍。

如何提高速度？我认为 TensorFlow 应该擅长完成这样的任务。

我应该将其构建为模型并使用预测来计算结果吗？这里没有什么可学的，所以我不知道该用什么。

Answer 1

如果你想从一个角度或者从一个角度的正弦和余弦构建 4x4 旋转矩阵，你可以这样做:

import tensorflow as tf

def make_rotation(alpha, axis):
    return make_rotation_sincos(tf.math.sin(alpha), tf.math.cos(alpha), axis)

def make_rotation_sincos(sin, cos, axis):
    axis = axis.strip().lower()
    zeros = tf.zeros_like(sin)
    ones = tf.ones_like(sin)
    if axis == 'x':
        rot = tf.stack([
            tf.stack([ ones, zeros, zeros], -1),
            tf.stack([zeros,   cos,  -sin], -1),
            tf.stack([zeros,   sin,   cos], -1),
        ], -2)
    elif axis == 'y':
        rot = tf.stack([
            tf.stack([  cos, zeros,   sin], -1),
            tf.stack([zeros,  ones, zeros], -1),
            tf.stack([ -sin, zeros,   cos], -1),
        ], -2)
    elif axis == 'z':
        rot = tf.stack([
            tf.stack([  cos,  -sin, zeros], -1),
            tf.stack([  sin,   cos, zeros], -1),
            tf.stack([zeros, zeros,  ones], -1),
        ], -2)
    else:
        raise ValueError('Invalid axis {!r}.'.format(axis))
    last_row = tf.expand_dims(tf.stack([zeros, zeros, zeros], -1), -2)
    last_col = tf.expand_dims(tf.stack([zeros, zeros, zeros, ones], -1), -1)
    return tf.concat([tf.concat([rot, last_row], -2), last_col], -1)

---

关于计算正向运动链，您可以使用tf.scan来完成。 。例如，假设初始形状 (?, 10):

# Make rotation matrices
rots = make_rotation(...)
rots_t = tf.transpose(rots, (1, 0, 2, 3))
out = tf.scan(tf.matmul, rots_t)[-1]