content: # Copyright (c) Meta Platforms, Inc. and affiliates.
  #
  # This source code is licensed under the MIT license found in the
  # LICENSE file in the root directory of this source tree.
  
  import os
  
  import pytest
  import torch
  
  import theseus as th
  
  NUM_DOFS = 7
  EE_NAME = "panda_virtual_ee_link"
  HOME_POSE = torch.Tensor([-0.1394, -0.0205, -0.0520, -2.0691, 0.0506, 2.0029, -0.9168])
  ERR_SCALING = torch.Tensor([1.0, 1.0, 1.0, 5.0, 5.0, 5.0])
  URDF_PATH = os.path.join(os.path.dirname(__file__), "data/panda_no_gripper.urdf")
  
  
  @pytest.fixture(params=[1, 3])
  def batch_size(request):
      return request.param
  
  
  @pytest.fixture
  def ee_pose_target(batch_size):
      torch.manual_seed(1)
  
      ee_pos_mid = torch.Tensor([0.6, 0.0, 0.5])
      ee_pos_range = torch.Tensor([0.1, 0.2, 0.2])
      ee_quat_mid = torch.Tensor([0.9383, 0.3442, -0.0072, -0.0318])
      ee_quat_range = torch.Tensor([0.5, 0.5, 0.5])
  
      ee_pos_dev = ee_pos_range * torch.randn(batch_size, 3)
      ee_rot_dev = ee_quat_range * torch.randn(batch_size, 3)
  
      ee_pos_target = ee_pos_mid + ee_pos_dev
      ee_quat_target = (
          th.SO3.unit_quaternion_to_SO3(
              ee_quat_mid / torch.linalg.norm(ee_quat_mid)
          ).compose(th.SO3().exp_map(ee_rot_dev))
      ).to_quaternion()
      ee_pose_target = th.SE3(
          x_y_z_quaternion=torch.cat([ee_pos_target, ee_quat_target], dim=-1),
          name="ee_pose_target",
      )
  
      return ee_pose_target
  
  
  @pytest.mark.parametrize("is_grad_enabled", [True, False])
  def test_ik_optimization(batch_size, ee_pose_target, is_grad_enabled):
      """Sets up inverse kinematics as an optimization problem that uses forward kinematics"""
  
      robot_model = th.eb.UrdfRobotModel(URDF_PATH, link_names=[EE_NAME])
  
      # Define cost (distance between desired and current ee pose)
      def ee_pose_err_fn(optim_vars, aux_vars):
          (theta,) = optim_vars
          (ee_pose_target,) = aux_vars
  
          ee_pose = robot_model.forward_kinematics(theta)[EE_NAME]
          pose_err = ee_pose_target.local(ee_pose)
  
          return pose_err
  
      # Set up optimization
      optim_vars = (th.Vector(NUM_DOFS, name="theta"),)
      aux_vars = (ee_pose_target,)
  
      cost_function = th.AutoDiffCostFunction(
          optim_vars,
          ee_pose_err_fn,
          6,
          aux_vars=aux_vars,
          name="ee_pose_err_fn",
          autograd_mode="dense",
      )
      objective = th.Objective()
      objective.add(cost_function)
      optimizer = th.LevenbergMarquardt(
          objective,
          max_iterations=25,
          step_size=0.5,
      )
      theseus_optim = th.TheseusLayer(optimizer)
  
      # Optimize
      theseus_inputs = {
          "theta": torch.tile(HOME_POSE.unsqueeze(0), (batch_size, 1)),
          "ee_pose_target": ee_pose_target,
      }
  
      with torch.set_grad_enabled(is_grad_enabled):
          updated_inputs, info = theseus_optim.forward(
              theseus_inputs,
              optimizer_kwargs={
                  "track_best_solution": True,
                  "verbose": True,
                  "track_error_history": True,
                  "damping": 0.1,
              },
          )
  
      # Check result
      torch.testing.assert_close(
          torch.zeros(batch_size), info.best_err, atol=1e-3, rtol=1e-3
      )
  