==========================
Inverse Dynamics Modeling
==========================

The foundational inverse dynamics model is based on all available data, i.e. it disregards the specific intance.
Generally speaking the model is trained as followed:
1. A wandb project manages this part of the project: https://wandb.ai/leon-gorissen/franka_wwl_demonstrator/sweeps
2. Sweeps in this project manage the hyperparameter tuning. I.e. wandb sweep server is employed
3. Agents get their hyperparameters from the sweep server. Training is carried out on our compute resources:
   a. If SWEEP_IP is exported, a new agent is registered to the sweep server.
      1. The training data is loaded from RDS using coscine python SDK.
      2. A new model is trained and evaluated.
      3. The model performance of the currently best performing model is retrieved from RDS using coscine python SDK.
      4. If the new model is better, it, the training hyperparameters and the sweep config are saved to RDS using coscine python SDK.
   b. Else a new sweep will be registered and this sweep_id will be used to register the agent to.
      * Training itself is the same. Hyperparameters can be retrieved from RDS if necessary, e.g. a crashed sweep server, to continue training in a meaningful area the hyperparameter space. 

Setup
=====

Please follow `official documentation <https://org.ngc.nvidia.com/setup/api-key>`_ on how to log into nvcr.io registry. Follow the instructions for docker.

Please log into `weights and biases <https://wandb.ai/home>`_ and retrieve your api key. Next we want to export the api key:

.. code-block:: bash

   export WANDB_API_KEY="your_api_key"

Next, export your username:

.. code-block:: bash

   export WANDB_ENTITY="your_username"

Next, export some notes, if you want:

.. code-block:: bash

   export WANDB_NOTES="your notes. here spaces are allowed"

Optional Setup:

.. code-block:: bash

   export CUDA_DEVICE="0"
   export NUM_MODELS="10"

Adding an Agent to an existing sweep
------------------------------------

To add an agent to an existing sweep, you need to retrieve the `sweep id <https://wandb.ai/leon-gorissen/franka_wwl_demonstrator/sweeps>`_ and export it:

.. code-block:: bash

   export WANDB_SWEEP_ID="your_sweep_id"

.. Note::
   Keep in mind, that a new sweep likely means, that the sweep configuration was changed. I.e. there are some more / other hyperparameters to be tuned. Thus, you might need to check which sweep your code supports.

Creating a new sweep config
---------------------------

If you want to create a new sweep, e.g. to investigate different architecuture, you need to revoke the SWEEP_IP variable: 

.. code-block:: bash

   unset WANDB_SWEEP_ID

Now you can adapt the source code in the dynamics_learning package in main.py.

Training the foundational model
===============================

To train the foundational model jump into the main git repo and:

.. code-block:: bash

   docker compose up -d dynamics_learning