From deaef459431adef9e9deea8b70568009b0e7f061 Mon Sep 17 00:00:00 2001
From: ebuehrle <43623224+ebuehrle@users.noreply.github.com>
Date: Tue, 14 Sep 2021 07:17:43 +0200
Subject: [PATCH] Add test for discriminator

---
 .../intersimple/gail/test_discriminator.py    | 33 ++++++++++++++++++-
 1 file changed, 32 insertions(+), 1 deletion(-)

diff --git a/scratch/etienne/intersimple/gail/test_discriminator.py b/scratch/etienne/intersimple/gail/test_discriminator.py
index 183cf27..1de614c 100644
--- a/scratch/etienne/intersimple/gail/test_discriminator.py
+++ b/scratch/etienne/intersimple/gail/test_discriminator.py
@@ -9,6 +9,37 @@ def test_image_concatenation():
     a = torch.tensor([[0.5]])
     sa = disc._concatenate(s, a)
 
+    assert s.shape == (1, 5, 200, 200)
+    assert a.shape == (1, 1)
     assert sa.shape == (1, 6, 200, 200)
     assert torch.allclose(sa[:, :5], 1.0 * s)
-    assert (sa[:, 5] == a).all()
+    assert (sa[:, 5] == a.unsqueeze(-1)).all()
+
+def test_image_concatenation3():
+    env = NRasterized()
+    disc = CnnDiscriminator(env)
+
+    s1 = env.reset()
+    a1 = 0.15
+    s2, _, _, _ = env.step(0.9)
+    a2 = 0.25
+    s3, _, _, _ = env.step(-0.9)
+    a3 = 0.35
+    
+    s = torch.stack([
+        torch.tensor(s1),
+        torch.tensor(s2),
+        torch.tensor(s3)
+    ], axis=0)
+    a = torch.tensor([
+        [a1],
+        [a2],
+        [a3],
+    ])
+    sa = disc._concatenate(s, a)
+
+    assert s.shape == (3, 5, 200, 200)
+    assert a.shape == (3, 1)
+    assert sa.shape == (3, 6, 200, 200)
+    assert torch.allclose(sa[:, :5], 1.0 * s)
+    assert (sa[:, 5] == a.unsqueeze(-1)).all()