diff --git a/flatland/envs/agent_chains.py b/flatland/envs/agent_chains.py
index 3e566ad0617a3c49ec69e1049be36231b705916f..e99b1dae3e02a08c333bf245a54ecb724881ca33 100644
--- a/flatland/envs/agent_chains.py
+++ b/flatland/envs/agent_chains.py
@@ -218,21 +218,21 @@ class MotionCheck(object):
         if "color" in dAttr:
             sColor = dAttr["color"]
             if sColor in [ "red", "purple" ]:
-                return (False, rcPos)
+                return False
 
         dSucc = self.G.succ[rcPos]
 
         # This should never happen - only the next cell of an agent has no successor
         if len(dSucc)==0:
             print(f"error condition - agent {iAgent} node {rcPos} has no successor")
-            return (False, rcPos)
+            return False
 
         # This agent has a successor
         rcNext = self.G.successors(rcPos).__next__()
         if rcNext == rcPos:  # the agent didn't want to move
-            return (False, rcNext)
+            return False
         # The agent wanted to move, and it can
-        return (True, rcNext)
+        return True
 
 
 
diff --git a/flatland/envs/agent_utils.py b/flatland/envs/agent_utils.py
index bf92637113a1c329f464e44c2c7ea2d01e996e28..23f8c032f304d3d41639d7f1b260c12a5575bbab 100644
--- a/flatland/envs/agent_utils.py
+++ b/flatland/envs/agent_utils.py
@@ -1,8 +1,4 @@
-from flatland.envs.rail_trainrun_data_structures import Waypoint
-import numpy as np
-
-from itertools import starmap
-from typing import Tuple, Optional, NamedTuple, List
+from typing import Tuple, Optional, NamedTuple
 
 from attr import attr, attrs, attrib, Factory
 
diff --git a/flatland/envs/rail_env.py b/flatland/envs/rail_env.py
index c8f75908aa7a5808acc8dbb1591afb9941bcfde1..75b10e8f71b45b078017f0e1ff9061249c594c75 100644
--- a/flatland/envs/rail_env.py
+++ b/flatland/envs/rail_env.py
@@ -549,7 +549,7 @@ class RailEnv(Environment):
             if agent.malfunction_handler.in_malfunction:
                 movement_allowed = False
             else:
-                movement_allowed, _ = self.motionCheck.check_motion(i_agent, agent.position) # TODO: Remove final_new_postion from motioncheck
+                movement_allowed = self.motionCheck.check_motion(i_agent, agent.position) # TODO: Remove final_new_postion from motioncheck
 
             if movement_allowed:
                 agent.position = agent_transition_data.position
@@ -567,7 +567,7 @@ class RailEnv(Environment):
                 agent.position = None
 
             ## Update rewards
-            # self.update_rewards(i_agent, agent, rail)
+            # self.update_rewards(i_agent, agent, rail) # TODO : Rewards - Fix this
 
             ## Update counters (malfunction and speed)
             agent.speed_counter.update_counter(agent.state)
@@ -577,7 +577,7 @@ class RailEnv(Environment):
             if agent.speed_counter.is_cell_entry:
                 agent.action_saver.clear_saved_action()
         
-        self.rewards_dict = {i_agent: 0 for i_agent in range(len(self.agents))} # TODO : Remove this
+        self.rewards_dict = {i_agent: 0 for i_agent in range(len(self.agents))} # TODO : Rewards - Remove this
         return self._get_observations(), self.rewards_dict, self.dones, info_dict # TODO : Will need changes?
 
     def record_timestep(self, dActions):