diff --git a/.gitignore b/.gitignore
index 235b179..1df75e9 100644
--- a/.gitignore
+++ b/.gitignore
@@ -3,4 +3,6 @@ notebooks/
 __pycache__/
 .ipynb_checkpoints/
 scripts/run_*.sh
-test*
\ No newline at end of file
+test*
+*frame2state*
+dmc2gym*
\ No newline at end of file
diff --git a/scripts/run.sh b/scripts/run.sh
index d44d7d2..41fe7b7 100755
--- a/scripts/run.sh
+++ b/scripts/run.sh
@@ -4,6 +4,6 @@ CUDA_VISIBLE_DEVICES=5 python train.py \
     --encoder_type pixel \
     --action_repeat 8 \
     --save_tb --pre_transform_image_size 100 --image_size 84 \
-    --work_dir ./tmp/cartpole \
+    --work_dir ./tmp/test \
     --agent curl_sac --frame_stack 3 \
     --seed -1 --critic_lr 1e-3 --actor_lr 1e-3 --eval_freq 10000 --batch_size 128 --num_train_steps 1000000 
\ No newline at end of file
diff --git a/train.py b/train.py
index 294dfa4..441fd14 100644
--- a/train.py
+++ b/train.py
@@ -172,7 +172,11 @@ def main():
     # stack several consecutive frames together
     if args.encoder_type == 'pixel':
         env = utils.FrameStack(env, k=args.frame_stack)
-    
+    else:
+        pos_dim = utils.get_pos_dim(args.domain_name,args.task_name)
+        env = utils.StateMask(env,pos_dim)
+       
+
     # make directory
     ts = time.gmtime() 
     ts = time.strftime("%m-%d", ts)    
diff --git a/utils.py b/utils.py
index ac2e150..fd1d9cb 100644
--- a/utils.py
+++ b/utils.py
@@ -8,6 +8,7 @@ import random
 from torch.utils.data import Dataset, DataLoader
 import time
 from skimage.util.shape import view_as_windows
+from dm_control import suite 
 
 class eval_mode(object):
     def __init__(self, *models):
@@ -198,6 +199,44 @@ class ReplayBuffer(Dataset):
     def __len__(self):
         return self.capacity 
 
+class StateMask(gym.Wrapper):
+    def __init__(self, env, pos_dim):
+        gym.Wrapper.__init__(self, env)
+        self.pos_dim = pos_dim
+        self.observation_space = gym.spaces.Box(
+            low=-np.inf,
+            high=np.inf,
+            shape=(pos_dim,),
+            dtype=env.observation_space.dtype
+        )
+
+    def reset(self):
+        obs = self.env.reset()
+        return obs[:self.pos_dim]
+       
+    def step(self, action):
+        obs, reward, done, info = self.env.step(action)
+        return obs[:self.pos_dim], reward, done, info
+
+def get_pos_dim(domain_name,task_name):
+    env = suite.load(domain_name,task_name)
+    ts = env.reset()
+    total_counts = 0
+    for k,v in ts.observation.items():
+        try: 
+            total_counts +=len(v)
+            if k == 'position' or k == 'orientations':
+                pos_counts = len(v)
+            #print(k,len(v))
+        except:
+            total_counts +=1
+            #print(k,1)
+            if k == 'position' or k == 'orientations':
+                pos_counts = 1
+            
+    return pos_counts
+
+
 class FrameStack(gym.Wrapper):
     def __init__(self, env, k):
         gym.Wrapper.__init__(self, env)