def generate_random_inverse_dynamics_data()

in differentiable_robot_model/data_utils.py [0:0]

• 37 lines of code
• 4 McCabe index (conditional complexity)

def generate_random_inverse_dynamics_data(robot_model, n_data):
    device = robot_model._device

    limits_per_joint = robot_model.get_joint_limits()
    joint_lower_bounds = np.asarray([joint["lower"] for joint in limits_per_joint])
    joint_upper_bounds = np.asarray([joint["upper"] for joint in limits_per_joint])
    joint_velocity_limits = 0.2 * np.asarray(
        [joint["velocity"] for joint in limits_per_joint]
    )
    q = torch.tensor(
        np.random.uniform(
            low=joint_lower_bounds, high=joint_upper_bounds, size=(n_data, 7)
        ),
        dtype=torch.float32,
    )
    qd = torch.tensor(
        np.random.uniform(
            low=-joint_velocity_limits, high=joint_velocity_limits, size=(n_data, 7)
        ),
        dtype=torch.float32,
    )
    qdd_des = torch.tensor(
        np.random.uniform(
            low=-joint_velocity_limits * 2.0,
            high=joint_velocity_limits * 2.0,
            size=(n_data, 7),
        ),
        dtype=torch.float32,
    )

    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 InverseDynamicsDataset(
        data={"q": q, "qd": qd, "qdd_des": qdd_des, "tau": torques}
    )