diff --git a/README.md b/README.md
index 1ed5680..2962c47 100644
--- a/README.md
+++ b/README.md
@@ -35,7 +35,7 @@ We welcome your contributions.
 
 If you use IntersectionZoo in your work, you are highly encouraged to cite our paper:
 
-V. Jayawardana, B. Freydt, A. Qu, C. Hickert, Z. Yan, C. Wu, "IntersectionZoo: Eco-driving for Benchmarking Multi-Agent Contextual Reinforcement Learning", 2024. 
+V. Jayawardana, B. Freydt, A. Qu, C. Hickert, Z. Yan, C. Wu, "IntersectionZoo: Eco-driving for Benchmarking Multi-Agent Contextual Reinforcement Learning", International Conference on Learning Representations (ICLR) 2025. 
 
 # Contributors
 
diff --git a/code/env/config.py b/code/env/config.py
index 769f9b3..4646b0e 100644
--- a/code/env/config.py
+++ b/code/env/config.py
@@ -39,7 +39,7 @@ class IntersectionZooEnvConfig(NamedTuple):
     logging: bool = True
     """ Verbose logging """
     visualize_sumo: bool = False
-    """ whether to use SUMO's GUI, doesn't always work very well with Ray"""
+    """ whether to use SUMO's GUI"""
     moves_output: Optional[Path] = None
     """ If specified exports vehicle trajectories in the given dir to be processed by MOVES to get accurate emissions"""
     trajectories_output: bool = False
diff --git a/code/env/environment.py b/code/env/environment.py
index c214ec3..0b9bb96 100644
--- a/code/env/environment.py
+++ b/code/env/environment.py
@@ -65,6 +65,7 @@ def __init__(self, config: EnvContext):
         self.uncontrolled_region_metrics = defaultdict(lambda: [])
         self.warmup_vehicles: Set[str] = set()
         self._agent_ids: Set[str] = {'mock'}
+        self.force_visualize = False
 
         self.action_space = (
             GymDict(
@@ -140,7 +141,7 @@ def reset(
         """
         super().reset(seed=seed, options=options)
         self.traci = start_sumo(
-            self.config,
+            self.config.update({"visualize_sumo": True}) if self.force_visualize else self.config,
             self.traci,
             self.prefix,
             self.get_task(),
@@ -1001,7 +1002,15 @@ def individual_reward(veh: Vehicle) -> float:
 
             speed = veh.speed
             threshold = self.config.threshold
-            emission = 0
+            emission = self.traffic_state.fuel_emissions_models.get_emissions_single_condition(
+                                        veh.emissions_type,
+                                        veh.speed,
+                                        veh.accel,
+                                        veh.slope,
+                                        self.task_context.temperature_humidity,
+                                        self.task_context.electric_or_regular,
+                                        is_rl(veh.id)
+            )
 
             if self.config.stop_penalty is not None and speed < threshold:
                 penalty += self.config.stop_penalty * (threshold - speed) / threshold
diff --git a/code/env_demo.py b/code/env_demo.py
index c0da7ae..ec133df 100644
--- a/code/env_demo.py
+++ b/code/env_demo.py
@@ -1,3 +1,25 @@
+# MIT License
+
+# Copyright (c) 2024 Vindula Jayawardana, Baptiste Freydt, Ao Qu, Cameron Hickert, Zhongxia Yan, Cathy Wu
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
 import argparse
 from pathlib import Path
 
@@ -69,4 +91,4 @@ def filter_rew(rew: dict):
     print("Reward:", filter_rew(reward))
 
 # Close the environment
-env.close()
\ No newline at end of file
+env.close()
diff --git a/code/policy_evaluation.py b/code/policy_evaluation.py
index fb53be6..3afa854 100644
--- a/code/policy_evaluation.py
+++ b/code/policy_evaluation.py
@@ -38,6 +38,7 @@
 
 parser.add_argument('--penetration', default=1.0, type=str, help='Eco drive adoption rate')
 parser.add_argument('--temperature_humidity', default='20_50', type=str, help='Temperature and humidity for evaluations')
+parser.add_argument('--visualize', default=False, type=bool, help='Visualize the agents in SUMO.')
 
 args = parser.parse_args()
 print(args)
@@ -70,10 +71,13 @@ def flatten_dict(d, parent_key='', sep='_'):
 
 for i, task in enumerate(tasks.list_tasks(False)):
     for _ in range(args.eval_per_task):
+        def edit_env(env):
+            env.force_visualize = args.visualize
+            env.set_task(task)
     
         algo.evaluation_workers.foreach_worker(
                 lambda ev: ev.foreach_env(
-                    lambda env: env.set_task(task)))
+                    lambda env: edit_env(env)))
         results = algo.evaluate()
 
         flattened_results = {**flatten_dict(results)}