pyrcf.components.agents.torchscript_agent_base
==============================================

.. py:module:: pyrcf.components.agents.torchscript_agent_base


Classes
-------

.. autoapisummary::

   pyrcf.components.agents.torchscript_agent_base.TorchScriptAgentBase


Module Contents
---------------

.. py:class:: TorchScriptAgentBase(model_file: str, input_dims: int, device: Literal['cpu', 0] = 'cpu', warmup_iterations: int = 10, default_kp: numpy.ndarray | float = 20, default_kd: numpy.ndarray | float = 1.0, update_rate: float = None, clock: pyrcf.utils.time_utils.ClockBase = PythonPerfClock(), dtype: torch.dtype = None)

   Bases: :py:obj:`pyrcf.components.agents.ml_agent_base.MLAgentBase`


   A (abstract) base class that provides functionalities to use a torchscript model file
   and use it's inference mode to perform control update.

   Child class has to override two methods:
       - update_input_to_model
       - update_cmd_from_model_output
   See docstrings for each method in MLAgentBase class.


   .. py:attribute:: _torch_device


   .. py:attribute:: _model


   .. py:attribute:: _device
      :value: 'cpu'



   .. py:attribute:: _dtype
      :value: None



   .. py:attribute:: _dims


   .. py:attribute:: _tensor_inp


   .. py:attribute:: _input_ndarray

      The input array to be updated by the child class using the current robot
      state and global plan commands. This will be passed to the model as input during
      inference


   .. py:attribute:: _latest_ctrl_cmd
      :type:  pyrcf.core.types.RobotCmd
      :value: None


      The RobotCmd object that can be use in the `update_input_to_model` method
      (if required) as the last sent command to the robot. This object has to be updated
      by the child class's `update_cmd_from_model_output` using the output from the model
      inference step.


   .. py:attribute:: _default_kp
      :value: 20



   .. py:attribute:: _default_kd
      :value: 1.0



   .. py:attribute:: _rate
      :value: None



   .. py:method:: _should_run()


   .. py:method:: initialise_robot_cmd(joint_states: pyrcf.core.types.JointStates)

      Override this method in child class if custom initilisation is required (e.g. joint name order).
      By default, this method sets the joint name and position values to be equal to the input joint states
      object, with kp and kd set to be the default (in constructor).



   .. py:method:: update_input_to_model(robot_state: pyrcf.core.types.RobotState, global_plan: pyrcf.core.types.GlobalMotionPlan, t: float, dt: float) -> numpy.ndarray
      :abstractmethod:


      Should update (`self._input_ndarray`) using appropriate values (input to model). This method
      has access to `self._latest_ctrl_cmd` (type `RobotCmd`) as well if needed.
      (NOTE: `self._latest_ctrl_cmd` is set to be the initial robot joint positions (zero velocities
      and efforts commands) with `default_kp` and `default_kd` at start).



   .. py:method:: update_cmd_from_model_output(model_output: numpy.ndarray, robot_state: pyrcf.core.types.RobotState, global_plan: pyrcf.core.types.GlobalMotionPlan, t: float, dt: float) -> None
      :abstractmethod:


      Should update `self._latest_robot_cmd` (type `RobotCmd`) using the output from the NN model.

      :param model_output: the numpy array created from the output tensor from the model
                           after the inference query was done. This is the output of the neural network. This
                           method should use this object to update `self._latest_robot_cmd` to be sent to the
                           robot.
      :type model_output: np.ndarray



   .. py:method:: get_action(robot_state: pyrcf.core.types.RobotState, global_plan: pyrcf.core.types.GlobalMotionPlan, t: float = None, dt: float = None) -> pyrcf.core.types.RobotCmd

      Should return the control command given the current robot state and global plan.

      :param robot_state: The current state information from the robot.
      :type robot_state: RobotState
      :param global_plan: the latest global plan generated by the global planner
                          used in the loop.
      :type global_plan: GlobalMotionPlan
      :param t: the current time signature of the control loop. Defaults to None
                (controllers may or may not need this).
      :type t: float, optional
      :param dt: the time since the last control loop. Defaults to None
                 (controllers may or may not need this).
      :type dt: float, optional

      :returns: The output control command to be sent to the robot.
      :rtype: RobotCmd



   .. py:method:: get_last_output() -> Tuple[pyrcf.core.types.LocalMotionPlan, pyrcf.core.types.RobotCmd]

      Should return the last computed outputs by this agent.

      :returns:

                Output local plan and control command from this agent.
                    NOTE: Either of these can be None, if the agent does not compute them.
      :rtype: Tuple[LocalMotionPlan, RobotCmd]