add tests

Author: reiniscimurs

tests/test_marl_world.yaml

@@ -0,0 +1,22 @@
+world:
+  height: 12  # the height of the world
+  width: 12   # the height of the world
+  step_time: 1.0
+  sample_time: 1.0
+  collision_mode: 'reactive'
+
+robot:
+  - number: 3
+    kinematics: {name: 'diff'}
+    distribution: {name: 'manual'}
+    shape: {name: 'circle', radius: 0.2}
+    vel_min: [ 0, -1.0 ]
+    vel_max: [ 1.0, 1.0 ]
+    state: [[3, 10, 0], [3, 6, 0], [3, 2, 0]]
+    goal: [[9, 9, 0], [8, 8, 0], [7, 7, 0]]
+    color: ['royalblue', 'red', 'green', 'orange', 'purple', 'yellow', 'cyan', 'magenta', 'lime', 'pink', 'brown']
+    arrive_mode: position
+    goal_threshold: 0.3
+
+    plot:
+      show_trajectory: False

tests/test_model.py

@@ -1,12 +1,16 @@
 from pathlib import Path
 
+import torch
+
+from robot_nav.SIM_ENV.marl_sim import MARL_SIM
 from robot_nav.models.RCPG.RCPG import RCPG
 from robot_nav.models.TD3.TD3 import TD3
 from robot_nav.models.CNNTD3.CNNTD3 import CNNTD3
 from robot_nav.models.SAC.SAC import SAC
 from robot_nav.models.DDPG.DDPG import DDPG
 from robot_nav.utils import get_buffer
 from robot_nav.SIM_ENV.sim import SIM
+from robot_nav.models.MARL.marlTD3 import TD3 as marlTD3
 import pytest
 
 PROJECT_ROOT = Path(__file__).resolve().parents[1]
@@ -91,3 +95,56 @@ def test_max_bound_models(model, state_dim):
         iterations=2,
         batch_size=8,
     )
+
+
+def test_marl_models():
+    sim = MARL_SIM("/tests/test_marl_world.yaml", disable_plotting=True)
+    test_model = marlTD3(
+        state_dim=11,
+        action_dim=2,
+        max_action=1,
+        num_robots=sim.num_robots,
+        device="cpu",
+        save_every=0,
+        load_model=False,
+    )  # instantiate a model
+
+    replay_buffer = get_buffer(
+        model=test_model,
+        sim=sim,
+        load_saved_buffer=False,
+        pretrain=False,
+        pretraining_iterations=0,
+        training_iterations=0,
+        batch_size=0,
+        buffer_size=100,
+    )
+
+    for _ in range(10):
+        connections = torch.tensor(
+            [[0.0 for _ in range(sim.num_robots - 1)] for _ in range(3)]
+        )
+        (
+            poses,
+            distance,
+            cos,
+            sin,
+            collision,
+            goal,
+            a,
+            reward,
+            positions,
+            goal_positions,
+        ) = sim.step([[0, 0] for _ in range(sim.num_robots)], connections)
+        state, terminal = test_model.prepare_state(
+            poses, distance, cos, sin, collision, a, goal_positions
+        )
+        replay_buffer.add(
+            state, [[0, 0] for _ in range(sim.num_robots)], reward, terminal, state
+        )
+
+    test_model.train(
+        replay_buffer=replay_buffer,
+        iterations=2,
+        batch_size=8,
+    )

tests/test_sim.py

@@ -1,7 +1,9 @@
 import os
 
 import pytest
+import torch
 
+from robot_nav.SIM_ENV.marl_sim import MARL_SIM
 from robot_nav.SIM_ENV.sim import SIM
 import numpy as np
 
@@ -12,7 +14,7 @@
 
 @skip_on_ci
 def test_sim():
-    sim = SIM("/tests/test_world.yaml")
+    sim = SIM("/tests/test_world.yaml", disable_plotting=True)
     robot_state = sim.env.get_robot_state()
     state = sim.step(1, 0)
     next_robot_state = sim.env.get_robot_state()
@@ -28,6 +30,28 @@ def test_sim():
     assert np.not_equal(robot_state[1], new_robot_state[1])
 
 
+def test_marl_sim():
+    sim = MARL_SIM("/tests/test_marl_world.yaml", disable_plotting=True)
+    robot_state = [sim.env.robot_list[i].state[:2] for i in range(3)]
+    connections = torch.tensor(
+        [[0.0 for _ in range(sim.num_robots - 1)] for _ in range(3)]
+    )
+
+    _ = sim.step([[1, 0], [1, 0], [1, 0]], connections)
+    next_robot_state = [sim.env.robot_list[i].state[:2] for i in range(3)]
+    for j in range(3):
+        assert np.isclose(robot_state[j][0], next_robot_state[j][0] - 1)
+        assert np.isclose(robot_state[j][1], robot_state[j][1])
+
+    assert len(sim.env.obstacle_list) == 0
+
+    sim.reset()
+    new_robot_state = [sim.env.robot_list[i].state[:2] for i in range(3)]
+    for j in range(3):
+        assert np.not_equal(robot_state[j][0], new_robot_state[j][0])
+        assert np.not_equal(robot_state[j][1], new_robot_state[j][1])
+
+
 @skip_on_ci
 def test_sincos():
     sim = SIM("/tests/test_world.yaml")