--- a
+++ b/tests/test.round2.py
@@ -0,0 +1,51 @@
+from osim.env import ProstheticsEnv, rect
+import numpy as np
+import unittest
+import math
+
+class SimulationTest(unittest.TestCase):
+    def test_reset(self):
+        env = ProstheticsEnv(visualize=False, difficulty=0)
+        o = env.reset()
+        self.assertEqual(type(o), list)
+        o = env.reset(project = False)
+        self.assertEqual(type(o), dict)
+        o = env.reset(project = True)
+        self.assertEqual(type(o), list)
+        
+        action = env.action_space.sample()
+        o,r,d,i = env.step(action, project = False)
+        self.assertRaises(KeyError, lambda : o["target_vel"])
+
+        env = ProstheticsEnv(visualize=False, difficulty=1)
+        env.reset()
+        o,r,d,i = env.step(action, project = False)
+        self.assertEqual(len(o["target_vel"]), 3)
+        
+        env.generate_new_targets(10)
+        for i in range(20):
+            o,r,d,i = env.step(action, project = False)
+
+        self.assertGreater(rect([2, 0])[0], 1.99)
+        self.assertLess(rect([2, math.pi/2.0])[0], 0.01)
+
+        env.reset()
+        env.generate_new_targets(10)
+        
+        # After 300 steps we should be far
+        self.assertGreater(np.sum( (env.targets[300,:] - np.array([1.25,0,0]))**2 ), 0.01)
+
+        state = env.osim_model.get_state()
+        env.osim_model.get_joint("ground_pelvis").get_coordinates(0).setSpeedValue(state, 5)
+        env.osim_model.set_state(state)
+        o1,r1,d,i = env.step(action, project = False)
+        env.osim_model.get_joint("ground_pelvis").get_coordinates(0).setSpeedValue(state, 1.25)
+        env.osim_model.set_state(state)
+        o2,r2,d,i = env.step(action, project = False)
+
+        self.assertGreater(o1["joint_vel"]["ground_pelvis"],o2["joint_vel"]["ground_pelvis"])
+        self.assertGreater(r2,r1)
+        
+
+if __name__ == '__main__':
+    unittest.main()