def generate_sine_motion_forward_dynamics_data()

in differentiable_robot_model/data_utils.py [0:0]

• 30 lines of code
• 1 McCabe index (conditional complexity)

def generate_sine_motion_forward_dynamics_data(robot_model, n_data, dt, freq):
    device = robot_model._device
    n_dofs = robot_model._n_dofs

    A = 0.7
    q_start = torch.zeros(robot_model._n_dofs)
    T = int(n_data * dt)
    t = torch.linspace(0.0, T - 1, n_data)
    pi = torch.acos(torch.zeros(1)).item() * 2

    q = A * torch.sin(2 * pi * freq * t).reshape(n_data, 1).repeat(1, n_dofs) + q_start
    qd = (
        2
        * pi
        * freq
        * A
        * torch.cos(2 * pi * freq * t).reshape(n_data, 1).repeat(1, n_dofs)
    )
    qdd_des = (
        -((2 * pi * freq) ** 2)
        * A
        * torch.sin(2 * pi * freq * t).reshape(n_data, 1).repeat(1, n_dofs)
    )

    q = q.to(device)
    qd = qd.to(device)
    qdd_des = qdd_des.to(device)
    torques = robot_model.compute_inverse_dynamics(
        q=q, qd=qd, qdd_des=qdd_des, include_gravity=True
    )

    return ForwardDynamicsDataset(
        data={"q": q, "qd": qd, "qdd_des": qdd_des, "tau": torques}
    )