<commit_before>import os
import numpy as np
import pydrake
from pydrake.solvers import ik

# Load our model from URDF
robot = pydrake.rbtree.RigidBodyTree(os.path.join(pydrake.getDrakePath(),
                                                  "examples/PR2/pr2.urdf"))

# Add a convenient frame, positioned 0.1m away from the r_gripper_palm_link
# along that link's x axis
robot.addFrame(
           pydrake.rbtree.RigidBodyFrame("r_hand_frame",
                                         robot.findLink("r_gripper_palm_link"),
                                         np.array([0.1, 0, 0]),
                                         np.array([0., 0, 0])))

# Make sure attribute access works on bodies
assert robot.world().name() == "world"

hand_frame_id = robot.findFrame("r_hand_frame").frame_index
base_body_id = robot.findLink('base_footprint').body_index

constraints = [
               # These three constraints ensure that the base of the robot is
               # at z = 0 and has no pitch or roll. Instead of directly
               # constraining orientation, we just require that the points at
               # [0, 0, 0], [1, 0, 0], and [0, 1, 0] in the robot's base's
               # frame must all be at z = 0 in world frame.
               # We don't care about the x or y position of the robot's base,
               # so we use NaN values to tell the IK solver not to apply a
               # constraint along those dimensions.
               ik.WorldPositionConstraint(robot, base_body_id,
                                          np.array([0.0, 0.0, 0.0]),
                                          np.array([np.nan, np.nan, 0.0]),
                                          np.array([np.nan, np.nan, 0.0])),
               ik.WorldPositionConstraint(robot, base_body_id,
                                          np.array([1.0, 0.0, 0.0]),
                                          np.array([np.nan, np.nan, 0.0]),
                                          np.array([np.nan, np.nan, 0.0])),
               ik.WorldPositionConstraint(robot, base_body_id,
                                          np.array([0.0, 1.0, 0.0]),
                                          np.array([np.nan, np.nan, 0.0]),
                                          np.array([np.nan, np.nan, 0.0])),

               # This constraint exactly specifies the desired position of the
               # hand frame we defined earlier.
               ik.WorldPositionConstraint(robot, hand_frame_id,
                                          np.array([0.0, 0.0, 0.0]),
                                          np.array([0.5, 0.0, 0.6]),
                                          np.array([0.5, 0.0, 0.6])),
               # And this specifies the orientation of that frame
               ik.WorldEulerConstraint(robot, hand_frame_id,
                                       np.array([0.0, 0.0, 0.0]),
                                       np.array([0.0, 0.0, 0.0]))
               ]

q_seed = robot.getZeroConfiguration()
options = ik.IKoptions(robot)
results = ik.inverseKinSimple(robot, q_seed, q_seed, constraints, options)

# results.INFO gives the output status of SNOPT. INFO = 1 is good, anything
# less than 10 is OK, and any info >= 10 indicates an infeasibility or failure
# of the optimizer.
assert results.INFO == 1

print repr(results.q_sol)
<commit_msg>more comment clarification<commit_after>import os
import numpy as np
import pydrake
from pydrake.solvers import ik

# Load our model from URDF
robot = pydrake.rbtree.RigidBodyTree(os.path.join(pydrake.getDrakePath(),
                                                  "examples/PR2/pr2.urdf"))

# Add a convenient frame, positioned 0.1m away from the r_gripper_palm_link
# along that link's x axis
robot.addFrame(
           pydrake.rbtree.RigidBodyFrame("r_hand_frame",
                                         robot.findLink("r_gripper_palm_link"),
                                         np.array([0.1, 0, 0]),
                                         np.array([0., 0, 0])))

# Make sure attribute access works on bodies
assert robot.world().name() == "world"

hand_frame_id = robot.findFrame("r_hand_frame").frame_index
base_body_id = robot.findLink('base_footprint').body_index

constraints = [
               # These three constraints ensure that the base of the robot is
               # at z = 0 and has no pitch or roll. Instead of directly
               # constraining orientation, we just require that the points at
               # [0, 0, 0], [1, 0, 0], and [0, 1, 0] in the robot's base's
               # frame must all be at z = 0 in world frame.
               # We don't care about the x or y position of the robot's base,
               # so we use NaN values to tell the IK solver not to apply a
               # constraint along those dimensions. This is equivalent to
               # placing a lower bound of -Inf and an upper bound of +Inf along
               # those axes.
               ik.WorldPositionConstraint(robot, base_body_id,
                                          np.array([0.0, 0.0, 0.0]),
                                          np.array([np.nan, np.nan, 0.0]),
                                          np.array([np.nan, np.nan, 0.0])),
               ik.WorldPositionConstraint(robot, base_body_id,
                                          np.array([1.0, 0.0, 0.0]),
                                          np.array([np.nan, np.nan, 0.0]),
                                          np.array([np.nan, np.nan, 0.0])),
               ik.WorldPositionConstraint(robot, base_body_id,
                                          np.array([0.0, 1.0, 0.0]),
                                          np.array([np.nan, np.nan, 0.0]),
                                          np.array([np.nan, np.nan, 0.0])),

               # This constraint exactly specifies the desired position of the
               # hand frame we defined earlier.
               ik.WorldPositionConstraint(robot, hand_frame_id,
                                          np.array([0.0, 0.0, 0.0]),
                                          np.array([0.5, 0.0, 0.6]),
                                          np.array([0.5, 0.0, 0.6])),
               # And this specifies the orientation of that frame
               ik.WorldEulerConstraint(robot, hand_frame_id,
                                       np.array([0.0, 0.0, 0.0]),
                                       np.array([0.0, 0.0, 0.0]))
               ]

q_seed = robot.getZeroConfiguration()
options = ik.IKoptions(robot)
results = ik.inverseKinSimple(robot, q_seed, q_seed, constraints, options)

# results.INFO gives the output status of SNOPT. INFO = 1 is good, anything
# less than 10 is OK, and any info >= 10 indicates an infeasibility or failure
# of the optimizer.
assert results.INFO == 1

print repr(results.q_sol)

{'fraction_non_alphanumeric': 0.09915303738317757, 'fraction_numerical': 0.02570093457943925, 'mean_word_length': 1.3584710743801653, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 68, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3740985', 'n_tokens_mistral': 1907, 'n_tokens_neox': 1784, 'n_words': 549}