From 5337897abb99a7a0c5cfe8d370ee20f85287159a Mon Sep 17 00:00:00 2001
From: Roman Vaxenburg <vaxenburg@gmail.com>
Date: Tue, 3 Dec 2024 19:11:50 -0500
Subject: [PATCH] Test switching to force actuators.

---
 tests/common.py         | 29 +++++++++++++++++++++++++++++
 tests/test_flywalker.py | 14 ++++++++++++++
 2 files changed, 43 insertions(+)
 create mode 100755 tests/common.py

diff --git a/tests/common.py b/tests/common.py
new file mode 100755
index 0000000..92eda9d
--- /dev/null
+++ b/tests/common.py
@@ -0,0 +1,29 @@
+"""Common functions and fixtures for flybody tests."""
+
+import numpy as np
+
+
+def is_force_actuator(physics, actuator_id=None):
+    """Check if actuator with given id is a force actuator with ctrlrange == (-1, 1).
+    If actuator_id is None, check all actuators.
+    """
+    physics.reset()
+    if actuator_id is None:
+        inds = [*range(physics.model.nu)]
+    else:
+        assert isinstance(actuator_id, int)
+        inds = [actuator_id]
+    for i in inds:
+        assert physics.model.actuator_gainprm[i][0] != 0.
+        assert np.all(physics.model.actuator_gainprm[i][1:] == 0.)
+        assert np.all(physics.model.actuator_biasprm[i][:] == 0.)
+        assert physics.model.actuator_gaintype[i] == 0.
+        assert physics.model.actuator_biastype[i] == 0.
+        if physics.model.actuator_trntype[i] != 5:
+            # Force actuator (on either joint or tendon).
+            assert np.all(physics.model.actuator_ctrlrange[i] == (-1, 1))
+        else:
+            # Adhesion actuator.
+            assert np.all(physics.model.actuator_ctrlrange[i] == (0, 1))
+
+    return True
diff --git a/tests/test_flywalker.py b/tests/test_flywalker.py
index 92c23db..e60c20e 100755
--- a/tests/test_flywalker.py
+++ b/tests/test_flywalker.py
@@ -7,6 +7,7 @@
 from dm_control.composer.observation.observable import base as observable_base
 
 from flybody.fruitfly.fruitfly import FruitFly
+from .common import is_force_actuator
 
 
 TEST_ACTION = 0.3561
@@ -120,6 +121,19 @@ def test_fly_bulletproof():
                     assert action_spec.maximum[i] == 1
 
 
+def test_force_actuators():
+    """Test switching to force actuators."""
+    fly = FruitFly(use_legs=True,
+                   use_wings=True,
+                   use_mouth=True,
+                   use_antennae=True,
+                   joint_filter=0.01,
+                   adhesion_filter=0.02,
+                   force_actuators=True)
+    physics = mjcf.Physics.from_mjcf_model(fly.mjcf_model)
+    assert is_force_actuator(physics)
+
+
 def test_filterexact():
     """Test `filterexact` actuator activation dynamics."""