Maniskill sweep

Author: Joseph Suarez

pufferlib/config/mani_skill.ini

@@ -1,19 +1,22 @@
 
 [base]
 package = mani_skill
-env_name = mani_pickcube mani_pushcube mani_peginsertion
+env_name = mani_pickcube mani_pushcube mani_stackcube mani_peginsertion
 policy_name = Policy
 rnn_name = Recurrent
 
 [env]
 num_envs = 4096
+sim_steps_per_control = 5
+control_freq = 100
+solver_position_iterations = 15
 
 [vec]
 backend = PufferEnv
 num_envs = 1
 
 [train]
-total_timesteps = 100_000_000
+total_timesteps = 15_000_000
 adam_beta1 = 0.9832254546070032
 adam_beta2 = 0.9996089758513379
 adam_eps = 0.0000024542110227211678
@@ -39,9 +42,28 @@ downsample = 0
 
 [sweep.train.total_timesteps]
 distribution = log_normal
-min = 5e6
-max = 15e6
-mean = 10e6
+min = 2e7
+max = 5e7
+mean = 4e7
 scale = time
 
+[sweep.env.sim_steps_per_control]
+distribution = int_uniform
+min = 1
+max = 10
+mean = 5
+scale = auto
 
+[sweep.env.control_freq]
+distribution = int_uniform
+min = 10
+max = 100
+mean = 20
+scale = auto
+
+[sweep.env.solver_position_iterations]
+distribution = int_uniform
+min = 4
+max = 30
+mean = 15
+scale = auto

pufferlib/environments/mani_skill/environment.py

@@ -13,23 +13,35 @@
 ALIASES = {
     'mani_pickcube': 'PickCube-v1',
     'mani_pushcube': 'PushCube-v1',
+    'mani_stackcube': 'StackCube-v1',
     'mani_peginsertion': 'PegInsertionSide-v1',
 }
 
-def env_creator(name='PickCube-v1'):
+def env_creator(name='PickCube-v1', **kwargs):
     return functools.partial(make, name)
 
-def make(name, num_envs=1, render_mode='rgb_array', buf=None, seed=0):
+def make(name, num_envs=1, render_mode='rgb_array', buf=None, seed=0, **kwargs):
     '''Create an environment by name'''
 
     if name in ALIASES:
         name = ALIASES[name]
 
-    return ManiPufferEnv(name, num_envs=num_envs, render_mode=render_mode, buf=buf, seed=seed)
+    return ManiPufferEnv(name, num_envs=num_envs, render_mode=render_mode, buf=buf, seed=seed, **kwargs)
 
 class ManiPufferEnv(pufferlib.PufferEnv):
-    def __init__(self, name, num_envs=1, render_mode='rgb_array', log_interval=16, buf=None, seed=0):
-        self.env = gym.make(name, reward_mode='delta', num_envs=num_envs, render_mode=render_mode)
+    def __init__(self, name, num_envs=1, solver_position_iterations=15,
+            sim_steps_per_control=5, control_freq=20, render_mode='rgb_array',
+            log_interval=16, buf=None, seed=0):
+        sim_freq = int(sim_steps_per_control * control_freq)
+        sim_config = {
+            'scene_config': {
+                'solver_position_iterations': solver_position_iterations
+            },
+            'sim_freq': sim_freq,
+            'control_freq': control_freq
+        }
+        self.env = gym.make(name, reward_mode='delta', num_envs=num_envs,
+            render_mode=render_mode, sim_config=sim_config)
         self.env = ManiSkillVectorEnv(self.env, auto_reset=True, ignore_terminations=False, record_metrics=True)
         self.agents_per_batch = num_envs
 
@@ -68,13 +80,20 @@ def _flatten_info(self, info):
 
     def reset(self, seed=0):
         obs, info = self.env.reset()
-        self.observations = obs
+        #self.observations = torch.nan_to_num(obs)
+        self.observations = torch.clamp(torch.nan_to_num(obs), -5, 5)
+        self.observations = obs / 20.0
         self._flatten_info(info)
         return obs, []
 
     def step(self, actions):
         obs, reward, terminated, truncated, info = self.env.step(actions)
-        self.observations = obs
+        collapsed = torch.where(torch.isnan(obs).sum(1) > 0)[0]
+        if len(collapsed) > 0:
+            obs, _ = self.env.reset(options={'env_idx': collapsed})
+
+        self.observations = torch.clamp(torch.nan_to_num(obs), -5, 5)
+        #self.observations = obs / 20.0 #torch.nan_to_num(obs)
         self.rewards = reward
         self.terminated = terminated
         self.truncated = truncated

pufferlib/models.py

@@ -182,6 +182,8 @@ def forward(self, observations, state):
         hidden = hidden.transpose(0, 1)
         #hidden = self.pre_layernorm(hidden)
         hidden, (lstm_h, lstm_c) = self.lstm.forward(hidden, lstm_state)
+        hidden = hidden.float()
+ 
         #hidden = self.post_layernorm(hidden)
         hidden = hidden.transpose(0, 1)
 

pufferlib/pufferl.py

@@ -863,8 +863,6 @@ def download(self):
 
 def train(env_name, args=None, vecenv=None, policy=None, logger=None):
     args = args or load_config(env_name)
-    vecenv = vecenv or load_env(env_name, args)
-    policy = policy or load_policy(args, vecenv)
 
     # Assume TorchRun DDP is used if LOCAL_RANK is set
     if 'LOCAL_RANK' in os.environ:
@@ -875,6 +873,12 @@ def train(env_name, args=None, vecenv=None, policy=None, logger=None):
         local_rank = int(os.environ["LOCAL_RANK"])
         print(f"rank: {local_rank}, MASTER_ADDR={master_addr}, MASTER_PORT={master_port}")
         torch.cuda.set_device(local_rank)
+        os.environ["CUDA_VISIBLE_DEVICES"] = str(local_rank)
+
+    vecenv = vecenv or load_env(env_name, args)
+    policy = policy or load_policy(args, vecenv)
+
+    if 'LOCAL_RANK' in os.environ:
         args['train']['device'] = torch.cuda.current_device()
         torch.distributed.init_process_group(backend='nccl', world_size=world_size)
         policy = policy.to(local_rank)
@@ -925,10 +929,12 @@ def train(env_name, args=None, vecenv=None, policy=None, logger=None):
 
 def eval(env_name, args=None, vecenv=None, policy=None):
     args = args or load_config(env_name)
-    args['vec'] = dict(backend='Serial', num_envs=1)
+    backend = args['vec']['backend']
+    if backend != 'PufferEnv':
+        backend = 'Serial'
+
+    args['vec'] = dict(backend=backend, num_envs=1)
     vecenv = vecenv or load_env(env_name, args)
-    if not isinstance(vecenv, pufferlib.vector.Serial):
-        raise pufferlib.APIUsageError('eval requires Serial vector env')
 
     policy = policy or load_policy(args, vecenv)
     ob, info = vecenv.reset()
@@ -954,11 +960,12 @@ def eval(env_name, args=None, vecenv=None, policy=None):
             print('\033[0;0H' + render + '\n')
             time.sleep(1/args['fps'])
         elif driver.render_mode == 'rgb_array':
-            import cv2
-            render = cv2.cvtColor(render, cv2.COLOR_RGB2BGR)
-            cv2.imshow('frame', render)
-            cv2.waitKey(1)
-            time.sleep(1/args['fps'])
+            pass
+            #import cv2
+            #render = cv2.cvtColor(render, cv2.COLOR_RGB2BGR)
+            #cv2.imshow('frame', render)
+            #cv2.waitKey(1)
+            #time.sleep(1/args['fps'])
 
         with torch.no_grad():
             ob = torch.as_tensor(ob).to(device)