Using BAM in MuJoCo (CPU)
=========================

This page explains how to plug BAM friction models into a standard MuJoCo
simulation running on CPU. The entry point is :class:`bam.mujoco.MujocoController`.

Installation
------------

BAM is available on PyPI. Install it with the ``mujoco`` extra to pull in the
MuJoCo dependency:

.. code-block:: bash

   pip install better-actuator-models[mujoco]

Overview
--------

At each simulation step, :class:`~bam.mujoco.MujocoController` does three things:

1. Optionally lowers the supply voltage by a drop proportional to the previous
   step's load, to model battery + cable resistance.
2. Computes the motor torque from a firmware-like P-controller — optionally
   clipping it to the firmware current limit — and applies it via
   ``mj_data.ctrl``.
3. Evaluates the BAM friction model and writes the result into
   ``mj_model.dof_frictionloss`` and ``mj_model.dof_damping``.

Loading a model
---------------

Use :func:`bam.model.load_model` to obtain a :class:`~bam.model.Model` object.
Two approaches are available.

**Bundled motor** — the library ships identified parameters for a set of common
servos:

.. code-block:: python

   from bam.model import load_model

   model = load_model(motor_name="xl330", model="m6")

Supported motor names: ``"xl320"``, ``"xl330"``, ``"mx106"``, ``"mx64"``,
``"erob80:50"``, ``"erob80:100"``.
Supported model variants: ``"m1"`` through ``"m6"`` (see :doc:`../theory/models`).

**Custom JSON** — parameters produced by your own identification run:

.. code-block:: python

   model = load_model("path/to/params.json")

XML setup
---------

Each actuator must be declared as a ``motor`` in the MJCF file (not
``position`` or ``velocity``). BAM overwrites ``frictionloss``, ``damping``,
and ``armature`` at runtime, so any value set in the XML will be ignored.

.. code-block:: xml

   <actuator>
     <motor name="joint_1" joint="joint_1" gear="1"/>
     ...
     <motor name="joint_n" joint="joint_n" gear="1"/>
   </actuator>


Instantiating the controller
-----------------------------

.. code-block:: python

   import mujoco
   from bam.mujoco import MujocoController

   mj_model = mujoco.MjModel.from_xml_file("robot.xml")
   mj_data  = mujoco.MjData(mj_model)

   controller = MujocoController(
      model=model,
      actuator=["joint_1", ..., "joint_n"],  # must match the motor name in the XML
      mujoco_model=mj_model,
      mujoco_data=mj_data,
   )

The ``actuator`` argument can take a single string or a list of strings, which 
allows the same motor model to drive multiple joints. Each string must 
match the ``name`` attribute of the ``<motor>``.

Simulation loop
---------------

Inside the loop, call :meth:`~bam.mujoco.MujocoController.set_q_target` to
provide the desired joint angle, then :meth:`~bam.mujoco.MujocoController.update`
before every ``mj_step``:

.. code-block:: python

   mujoco.mj_resetData(mj_model, mj_data)
   controller.reset(mj_data.qpos)

   joint_names = ["joint_1", ..., "joint_n"]
   target_angles = [...]

   while True:
      for joint_name, target_angle in zip(joint_names, target_angles):
         controller.set_q_target(joint_name, target_angle)
      controller.update()
      mujoco.mj_step(mj_model, mj_data)

:meth:`~bam.mujoco.MujocoController.reset` should be called after every
``mj_resetData`` to clear the internal velocity and torque state.

Voltage drop (optional)
-----------------------

Real batteries and cables introduce a voltage drop proportional to the total
current draw. BAM models this as:

.. math::

   V_\text{eff} = V_\text{in} - g_\text{drop} \sum_i |\tau_i|

where :math:`g_\text{drop}` is ``vin_drop_gain`` (approximately
:math:`R / K_t`) and the sum runs over all controlled joints.
A hard lower bound ``vin_min`` can be set to prevent the effective voltage
from collapsing under heavy load:

.. code-block:: python

   controller = MujocoController(
      model=model,
      actuator=["joint_1", ..., "joint_n"],
      mujoco_model=mj_model,
      mujoco_data=mj_data,
      vin_drop_gain=0.5,   # [V/Nm]
      vin_min=6.0,         # [V]
   )

Current clipping (optional)
---------------------------

Servo firmwares can cap the motor current to protect the hardware. BAM reproduces
this saturation with the ``max_current`` parameter: the motor current
:math:`I = \tau / K_t` is clipped to ``[-max_current, max_current]``, which is
equivalent to clipping the motor torque to :math:`\pm\,\texttt{max\_current}\cdot K_t`.

.. code-block:: python

   controller = MujocoController(
      model=model,
      actuator=["joint_1", ..., "joint_n"],
      mujoco_model=mj_model,
      mujoco_data=mj_data,
      max_current=1.75,   # firmware current limit [A]
   )

Leave it at ``None`` (default) to disable current clipping.

Multi-actuator config file
--------------------------

For robots with many joints, :func:`bam.mujoco.load_config` loads a JSON
configuration file that maps each group of joints to a model:

.. code-block:: python

   from bam.mujoco import load_config

   controllers, dof_to_controller = load_config(
      path="config.json",
      mujoco_model=mj_model,
      mujoco_data=mj_data,
      kp=125.0,
      vin=7.5,
   )

The config file has the following structure:

.. code-block:: json

   {
      "arm": {
         "dofs": ["shoulder", "elbow"],
         "model": {
            "kt": 1.6224667906987444,
            "R": 3.949433673232461,
            "armature": 0.011951238325312509,
            "friction_base": 0.09038677246291783,
            "friction_viscous": 0.011691602145974832,
            "model": "m1",
            "actuator": "mx64"
         },
         "error_gain": 1.0,
         "max_pwm": 885
      },
      "leg": {
         "dofs": ["hip", "knee", "ankle"],
         "model": {
            "kt": 2.1913757006745245,
            "R": 2.9649903987776804,
            "armature": 0.026609234235148084,
            "friction_base": 0.10352026623606064,
            "friction_viscous": 0.03520238029013507,
            "model": "m1",
            "actuator": "mx106"
         },
         "error_gain": 1.0,
         "max_pwm": 885
      }
   }

``controllers`` is a dict keyed by group name; ``dof_to_controller`` maps each
DOF name back to its group.

API reference
-------------

- :class:`bam.mujoco.MujocoController`
- :func:`bam.mujoco.load_config`
- :func:`bam.model.load_model`