i think i did the converstion correct

wrapper/reciever_econ_wrapper.py

@@ -6,16 +6,15 @@ from ai_economist.foundation.base import base_env
 import gym
 import gym.spaces
 import numpy as np
-from  base_econ_wrapper import BaseEconVecEnv
 
 from stable_baselines3.common.vec_env.base_vec_env import VecEnv, VecEnvIndices, VecEnvObs, VecEnvStepReturn
 from stable_baselines3.common.vec_env.util import copy_obs_dict, dict_to_obs, obs_space_info
-
+from wrapper.base_econ_wrapper import BaseEconWrapper 
 from ai_economist import foundation
 
-class RecieverEconVecEnv(gym.Env):
+class RecieverEconWrapper(gym.Env):
     """Reciever part of BaseEconVecEnv. Filters by agent class and presents gym api to RL algos. Enables multi threading learning for different agent types."""
-    def __init__(self, base_econ: BaseEconVecEnv, agent_classname: str):
+    def __init__(self, base_econ: BaseEconWrapper, agent_classname: str):
         self.base_econ=base_econ
         base_econ.register_vote()
         self.econ=base_econ.env
@@ -23,7 +22,7 @@ class RecieverEconVecEnv(gym.Env):
         self.agnet_idx=list(self.econ.world._agent_class_idx_map[agent_classname])
         self.idx_to_index={}
         #create idx to index map
-        for i in range(len(self.agnet_idx)):
+        for i in range(len(self.agnet_idx)): 
             self.idx_to_index[self.agnet_idx[i]]=i
         first_idx=self.agnet_idx[0]
     
@@ -36,13 +35,16 @@ class RecieverEconVecEnv(gym.Env):
     def _dict_idx_to_index(self, data):
         data_out={}
         for k,v in data.items():
-            data_out[self.idx_to_index[k]]=v
+            if k in self.idx_to_index:
+                index=self.idx_to_index[k]
+                data_out[index]=v
         return data_out
 
     def _dict_index_to_idx(self, data):
         data_out={}
         for k,v in data.items():
-            data_out[self.agnet_idx[k]]=v
+            idx=self.agnet_idx[k]
+            data_out[idx]=v
         return data_out
 
     def step_wait(self):