.. Comment: this file is automatically generated by `update_example_docs.py`.
   It should not be modified manually.

Trajectory Optimization
==========================================


Basic Trajectory Optimization using PyRoKi.

Robot going over a wall, while avoiding world-collisions.

All examples can be run by first cloning the PyRoki repository, which includes the ``pyroki_snippets`` implementation details.



.. code-block:: python
        :linenos:


        import time
        from typing import Literal

        import numpy as np
        import pyroki as pk
        import trimesh
        import tyro
        import viser
        from viser.extras import ViserUrdf
        from robot_descriptions.loaders.yourdfpy import load_robot_description

        import pyroki_snippets as pks


        def main(robot_name: Literal["ur5", "panda"] = "panda"):
            if robot_name == "ur5":
                urdf = load_robot_description("ur5_description")
                down_wxyz = np.array([0.707, 0, 0.707, 0])
                target_link_name = "ee_link"

                # For UR5 it's important to initialize the robot in a safe configuration;
                # the zero-configuration puts the robot aligned with the wall obstacle.
                default_cfg = np.zeros(6)
                default_cfg[1] = -1.308
                robot = pk.Robot.from_urdf(urdf, default_joint_cfg=default_cfg)

            elif robot_name == "panda":
                urdf = load_robot_description("panda_description")
                target_link_name = "panda_hand"
                down_wxyz = np.array([0, 0, 1, 0])  # for panda!
                robot = pk.Robot.from_urdf(urdf)

            else:
                raise ValueError(f"Invalid robot: {robot_name}")

            robot_coll = pk.collision.RobotCollision.from_urdf(urdf)

            # Define the trajectory problem:
            # - number of timesteps, timestep size
            timesteps, dt = 50, 0.02
            # - the start and end poses.
            start_pos, end_pos = np.array([0.5, -0.3, 0.2]), np.array([0.5, 0.3, 0.2])

            # Define the obstacles:
            # - Ground
            # ground_coll = pk.collision.HalfSpace.from_point_and_normal(
            #     np.array([0.0, 0.0, 0.0]), np.array([0.0, 0.0, 1.0])
            # )
            # - Wall (using Box collision geometry)
            wall_height = 0.4
            wall_width = 0.1
            wall_length = 0.4
            wall_coll = pk.collision.Box.from_extent(
                extent=np.array([wall_length, wall_width, wall_height]),
                position=np.array([0.5, 0.0, wall_height / 2]),
            )
            # world_coll = [ground_coll, wall_coll]

            # TODO: Constraints don't work with ground collision at the moment,
            # because the robot is in collision already with it (which destabilizes things).
            # We will fix this if we can have better robot collision geometries / handling.
            world_coll = [wall_coll]

            traj = pks.solve_trajopt(
                robot,
                robot_coll,
                world_coll,
                target_link_name,
                start_pos,
                down_wxyz,
                end_pos,
                down_wxyz,
                timesteps,
                dt,
            )
            traj = np.array(traj)

            # Visualize!
            server = viser.ViserServer()
            urdf_vis = ViserUrdf(server, urdf)
            server.scene.add_grid("/grid", width=2, height=2, cell_size=0.1)
            server.scene.add_mesh_trimesh(
                "wall_box",
                trimesh.creation.box(
                    extents=(wall_length, wall_width, wall_height),
                    transform=trimesh.transformations.translation_matrix(
                        np.array([0.5, 0.0, wall_height / 2])
                    ),
                ),
            )
            for name, pos in zip(["start", "end"], [start_pos, end_pos]):
                server.scene.add_frame(
                    f"/{name}",
                    position=pos,
                    wxyz=down_wxyz,
                    axes_length=0.05,
                    axes_radius=0.01,
                )

            slider = server.gui.add_slider(
                "Timestep", min=0, max=timesteps - 1, step=1, initial_value=0
            )
            playing = server.gui.add_checkbox("Playing", initial_value=True)

            while True:
                if playing.value:
                    slider.value = (slider.value + 1) % timesteps

                urdf_vis.update_cfg(traj[slider.value])
                time.sleep(1.0 / 10.0)


        if __name__ == "__main__":
            tyro.cli(main)