Add restitution in manual collision detection

tests/test_envs.py

@@ -1,4 +1,5 @@
 import gym
+import numpy as np
 import pytest
 
 import mujoco_maze
@@ -6,15 +7,43 @@ import mujoco_maze
 
 @pytest.mark.parametrize("maze_id", mujoco_maze.TaskRegistry.keys())
 def test_ant_maze(maze_id):
-    env = gym.make("Ant{}-v0".format(maze_id))
-    assert env.reset().shape == (30,)
-    s, _, _, _ = env.step(env.action_space.sample())
-    assert s.shape == (30,)
+    for i in range(2):
+        env = gym.make(f"Ant{maze_id}-v{i}")
+        assert env.reset().shape == (30,)
+        s, _, _, _ = env.step(env.action_space.sample())
+        assert s.shape == (30,)
 
 
 @pytest.mark.parametrize("maze_id", mujoco_maze.TaskRegistry.keys())
 def test_point_maze(maze_id):
-    env = gym.make("Point{}-v0".format(maze_id))
-    assert env.reset().shape == (7,)
-    s, _, _, _ = env.step(env.action_space.sample())
-    assert s.shape == (7,)
+    for i in range(2):
+        env = gym.make(f"Point{maze_id}-v{i}")
+        assert env.reset().shape == (7,)
+        s, _, _, _ = env.step(env.action_space.sample())
+        assert s.shape == (7,)
+
+
+@pytest.mark.parametrize("maze_id", mujoco_maze.TaskRegistry.keys())
+def test_collision_lines(maze_id):
+    env = gym.make(f"Point{maze_id}-v0")
+    if maze_id == "UMaze":
+        assert len(env.unwrapped._collision.lines) == 16
+    structure = env.unwrapped._maze_structure
+    scaling = env.unwrapped._maze_size_scaling
+    init_x = env.unwrapped._init_torso_x
+    init_y = env.unwrapped._init_torso_y
+
+    def check_pos(pos):
+        x_orig = (pos.real + init_x) / scaling
+        y_orig = (pos.imag + init_y) / scaling
+        return structure[int(round(y_orig))][int(round(x_orig))]
+
+    for line in env.unwrapped._collision.lines:
+        mid = (line.p1 + line.p2) / 2
+        p2p1 = line.p2 - line.p1
+        cell1 = check_pos(mid + 0.1 * p2p1 * np.complex(0.0, -1.0))
+        cell2 = check_pos(mid + 0.1 * p2p1 * np.complex(0.0, 1.0))
+        if cell1.is_block():
+            assert not cell2.is_block()
+        else:
+            assert cell2.is_block()

tests/test_intersect.py

@@ -0,0 +1,70 @@
+import numpy as np
+import pytest
+
+from mujoco_maze.maze_env_utils import Line
+
+
+@pytest.mark.parametrize(
+    "l1p1, l1p2, l2p1, l2p2, none",
+    [
+        ((0.0, 0.0), (1.0, 0.0), (0.0, -1.0), (1.0, 1.0), False),
+        ((1.0, 1.0), (2.0, 3.0), (-1.0, 1.5), (1.5, 1.0), False),
+        ((1.5, 1.5), (2.0, 3.0), (-1.0, 1.5), (1.5, 1.0), True),
+        ((0.0, 0.0), (2.0, 0.0), (1.0, 0.0), (1.0, 3.0), False),
+    ],
+)
+def test_intersect(l1p1, l1p2, l2p1, l2p2, none):
+    l1 = Line(l1p1, l1p2)
+    l2 = Line(l2p1, l2p2)
+    i1 = l1.intersect(l2)
+    i2 = line_intersect(l1p1, l1p2, l2p1, l2p2)
+    if none:
+        assert i1 is None and i2 is None
+    else:
+        assert i1 is not None
+        np.testing.assert_array_almost_equal(i1, np.array(i2))
+
+
+def line_intersect(pt1, pt2, ptA, ptB):
+    """
+    Taken from https://www.cs.hmc.edu/ACM/lectures/intersections.html
+    Returns the intersection of Line(pt1,pt2) and Line(ptA,ptB).
+    """
+    import math
+
+    DET_TOLERANCE = 0.00000001
+
+    # the first line is pt1 + r*(pt2-pt1)
+    # in component form:
+    x1, y1 = pt1
+    x2, y2 = pt2
+    dx1 = x2 - x1
+    dy1 = y2 - y1
+
+    # the second line is ptA + s*(ptB-ptA)
+    x, y = ptA
+    xB, yB = ptB
+    dx = xB - x
+    dy = yB - y
+
+    DET = -dx1 * dy + dy1 * dx
+
+    if math.fabs(DET) < DET_TOLERANCE:
+        return None
+
+    # now, the determinant should be OK
+    DETinv = 1.0 / DET
+
+    # find the scalar amount along the "self" segment
+    r = DETinv * (-dy * (x - x1) + dx * (y - y1))
+
+    # find the scalar amount along the input line
+    s = DETinv * (-dy1 * (x - x1) + dx1 * (y - y1))
+
+    # return the average of the two descriptions
+    xi = (x1 + r * dx1 + x + s * dx) / 2.0
+    yi = (y1 + r * dy1 + y + s * dy) / 2.0
+    if r >= 0 and 0 <= s <= 1:
+        return xi, yi
+    else:
+        return None