Model predicting robot offsets not learning

Hi,

I am currently trying to predict a serial robot’s offsets by using a neural network on Pytorch. The method is the following :

  • Sampling random a vector of joint_angles.
  • Feeding neural network with it in order to predict corresponding offsets.
  • Computing FK(joint_angles+offsets) to get estimated_effector_position corresponding to the estimated joints states.
  • Moving the robot with joint_angles and measure real_effector_position which contains informations about the real offsets
  • Compute MSE(real_effector_position,estimated_effector_position) and then backprop error.
    The network should output offsets that minimize the error with training.

Formally, I’m trying to solve the following optimization problem :
arg min(offsets) ||FK(joint_angles+offsets) - real_effector_position||

I’m using pytroch to implement this but the network doesn’t learn and the loss stay the same. Here is my code

robot = BraccioRobot()
loss = 0
model = OffsetMLP(3,64)
opti = optim.Adam(net.parameters(), lr=3e-4)
for i in range(10000): 
    desired_joints_state = torch.FloatTensor(np.random.randint(0,180,(10,3))) # Get joint_angles
    offset = model(desired_joints_state/180)  * 15                            # Predict offset
    angles = desired_joints_state + offset                                    # Effective angles
    x,y = robot.forward_kinematic(((angles-90)*np.pi)/180)                    # Compute FK with offsets
    x_t,y_t = robot.offset_forward_kinematic(((desired_joints_state-90)*np.pi)/180) # Real joints_states
    loss = F.mse_loss(x,x_t) + F.mse_loss(y,y_t)                              # Compute loss
    print(loss)
    opti.zero_grad()
    loss.backward()
    opti.step()

the robot class :

class BraccioRobot():
    def __init__(self,) :
        self.proportions = np.array([0.17,0.29,0.285,0.25])
        self.joints_offsets = np.array([10,9,8]) * (np.pi/180.0)
    
    def forward_kinematic(self,joint_states):
        # Compute standard FK
        parent_angle = 0
        x = torch.zeros(joint_states.shape[0])
        y = torch.zeros(joint_states.shape[0])
        for i in range(joint_states.shape[1]):
            x += self.proportions[i+1] * torch.sin(joint_states[:,i] + parent_angle)
            y += self.proportions[i+1] * torch.cos(joint_states[:,i] + parent_angle)
            parent_angle += joint_states[:,i]
        y += self.proportions[0]
        return x,y # Effector position 

    def offset_forward_kinematic(self,joint_states):
        # Compute FK with added offsets
        parent_angle = 0
        x = torch.zeros(joint_states.shape[0])
        y = torch.zeros(joint_states.shape[0])
        for i in range(joint_states.shape[1]):
            x += self.proportions[i+1] * torch.sin(joint_states[:,i] + parent_angle + self.joints_offsets[i])
            y += self.proportions[i+1] * torch.cos(joint_states[:,i] + parent_angle + self.joints_offsets[i])
            parent_angle += joint_states[:,i]
        y += self.proportions[0]
        return x,y # Effector position

and the model :

class OffsetMLP(nn.Module):

    def __init__(self,n_joints,n_hidden):
        super(OffsetMLP,self).__init__()
        self.fc1 = nn.Linear(n_joints,n_hidden)
        self.fc2 = nn.Linear(n_hidden,n_joints)
        self.outputAct = nn.Tanh()

    def forward(self,desired_joints_state):
        offset = self.fc1(desired_joints_state)
        offset = F.relu(offset)
        offset = self.fc2(offset)
        offset = self.outputAct(offset)
        return offset

the offset_forward_kinematic function is used to emulate real robot movement. It just add the offsets that the network must learn. In this case it is 3 simple constant. I’m guessing the error is somewhere in the forward_kinematic function with the gradient maybe not propagating as it should but I can’t seem to find a solution. Thanks in advance for any fix.