reformat and code clean up

dedo/__init__.py

@@ -44,4 +44,4 @@
                  entry_point='dedo.envs:'+cls_nm, order_enforce=False)
 
 # Register dual-arm robot tasks.
-register(id='HangGarmentRobot-v1', entry_point='dedo.envs:DeformRobotEnv')
+register(id='HangGarmentRobot-v1', entry_point='dedo.envs:DeformRobotEnv', order_enforce=False)

dedo/demo.py

@@ -14,6 +14,7 @@
 """
 import matplotlib.pyplot as plt
 from matplotlib import interactive
+
 interactive(True)
 import numpy as np
 
@@ -56,13 +57,13 @@ def play(env, num_episodes, args):
         step = 0
         # input('Reset done; press enter to start episode')
         while True:
-            assert(not isinstance(env.action_space, gym.spaces.Discrete))
+            assert (not isinstance(env.action_space, gym.spaces.Discrete))
             print('step', step)
             act = env.action_space.sample()  # in [-1,1]
-            noise_act = 0.1*act
+            noise_act = 0.1 * act
             act = policy_simple(obs, noise_act, args.task, step)
             next_obs, rwd, done, info = env.step(act)
-            if args.viz and (args.cam_resolution > 0) and step%10 == 0:
+            if args.viz and (args.cam_resolution > 0) and step % 10 == 0:
                 plt.imshow(next_obs)
             if done:
                 break
@@ -73,7 +74,7 @@ def play(env, num_episodes, args):
 
 
 def main(args):
-    assert('Robot' not in args.env), 'This is a simple demo for anchors only'
+    assert ('Robot' not in args.env), 'This is a simple demo for anchors only'
     np.set_printoptions(precision=4, linewidth=150, suppress=True)
     kwargs = {'args': args}
     env = gym.make(args.env, **kwargs)

dedo/demo_preset.py

@@ -17,10 +17,10 @@
 import gym
 import scipy.linalg
 from matplotlib import interactive
+
 interactive(True)
 import numpy as np
 
-
 from dedo.utils.args import get_args
 from dedo.utils.anchor_utils import create_anchor_geom
 from dedo.utils.preset_info import preset_traj
@@ -29,6 +29,7 @@
 
 from dedo.utils.bullet_manipulator import convert_all
 
+
 def play(env, num_episodes, args):
     if args.task == 'ButtonProc':
         deform_obj = 'cloth/button_cloth.obj'
@@ -93,11 +94,11 @@ def play(env, num_episodes, args):
         traj_ori = preset_wp.get('a_theta', None)
         if traj_ori is not None:
             traj_ori = convert_all(np.array(traj_ori), 'theta_to_sin_cos')
-            n_repeats = traj.shape[0]//len(traj_ori)
+            n_repeats = traj.shape[0] // len(traj_ori)
             traj_ori = np.repeat(traj_ori, n_repeats, axis=0)
             print('traj_ori', traj_ori.shape, traj_ori)
-            assert(traj_ori.shape[0] == traj.shape[0])
-            assert(traj_ori.shape[1] == 6)  # Euler sin,sin,sin,cos,cos,cos
+            assert (traj_ori.shape[0] == traj.shape[0])
+            assert (traj_ori.shape[1] == 6)  # Euler sin,sin,sin,cos,cos,cos
             traj = np.hstack([traj, traj_ori])
 
         gif_frames = []
@@ -112,9 +113,9 @@ def play(env, num_episodes, args):
             next_obs, rwd, done, info = env.step(act, unscaled=True)
             rwds.append(rwd)
 
-            if done and vidwriter is not None: # Record the internal steps after anchor drop
+            if done and vidwriter is not None:  # Record the internal steps after anchor drop
                 for ob in info['final_obs'][1:]:
-                    vidwriter.write(np.uint8(ob[...,::-1]*255))
+                    vidwriter.write(np.uint8(ob[..., ::-1] * 255))
             if args.cam_resolution > 0:
                 img = env.render(mode='rgb_array', width=args.cam_resolution,
                                  height=args.cam_resolution)
@@ -134,7 +135,7 @@ def play(env, num_episodes, args):
         if args.use_wandb:
             mean_rwd = np.sum(rwds)
             for i in range(31):
-                wandb.log({'rollout/ep_rew_mean': mean_rwd, 'Step':i}, step=i)
+                wandb.log({'rollout/ep_rew_mean': mean_rwd, 'Step': i}, step=i)
         if vidwriter is not None:
             vidwriter.release()
 
@@ -159,7 +160,7 @@ def merge_traj(traj_a, traj_b):
 
 def build_traj(env, preset_wp, left_or_right, anchor_idx, ctrl_freq, robot):
     if robot is not None:
-        init_anc_pos = env.robot.get_ee_pos(left=anchor_idx>0)
+        init_anc_pos = env.robot.get_ee_pos(left=anchor_idx > 0)
     else:
         anc_id = list(env.anchors.keys())[anchor_idx]
         init_anc_pos = env.anchors[anc_id]['pos']
@@ -171,15 +172,14 @@ def build_traj(env, preset_wp, left_or_right, anchor_idx, ctrl_freq, robot):
     # exit(1)
     # WARNING: old code below.
 
-
     from scipy.interpolate import interp1d
     wpt = np.concatenate([[init_anc_pos], wp[:, :3]], axis=0)
     ids = np.arange(wpt.shape[0])
     interp_type = 'linear'
     # Creates the respective time interval for each way point
     interp_i = []
     for i, num_step in enumerate(steps):
-        interp_i.append(np.linspace(i, i+1, num_step, endpoint=False))
+        interp_i.append(np.linspace(i, i + 1, num_step, endpoint=False))
 
     interp_i = np.concatenate(interp_i)
     # interp_i = np.linspace(0, 1, steps[0], endpoint=False) # np.arange(0, wpt.shape[0]-1, 0.1)
@@ -189,24 +189,22 @@ def build_traj(env, preset_wp, left_or_right, anchor_idx, ctrl_freq, robot):
 
     traj = np.array([xi, yi, zi]).T
 
-
     dv = (traj[1:] - traj[:-1])  # * ctrl_freq
 
-
     # Calculating the avg velocity for each control step
     chunks = []
     chunk_size = int(np.round(ctrl_freq))
     start = 0
     for i in range(99999):
 
-        if start+chunk_size > dv.shape[0]:
+        if start + chunk_size > dv.shape[0]:
             # last chunk
             chunk_size = dv.shape[0] - start
-        chunk = dv[start:start+chunk_size]
+        chunk = dv[start:start + chunk_size]
         mean_chunk = np.sum(chunk, axis=0, keepdims=True)
-        mean_chunk = np.repeat(mean_chunk, chunk_size, axis=0, ) # scale back to original shape
+        mean_chunk = np.repeat(mean_chunk, chunk_size, axis=0, )  # scale back to original shape
         chunks.append(mean_chunk)
-        start = start+chunk_size
+        start = start + chunk_size
         if start >= dv.shape[0]:
             break
 

dedo/envs/deform_robot_env.py

@@ -26,7 +26,7 @@
 
 
 class DeformRobotEnv(DeformEnv):
-    ORI_SIZE = 3*2    # 3D position + sin,cos for 3 Euler angles
+    ORI_SIZE = 3 * 2  # 3D position + sin,cos for 3 Euler angles
     FING_DIST = 0.01  # default finger distance
 
     def __init__(self, args):
@@ -44,7 +44,7 @@ def __init__(self, args):
     def unscale_pos(act, unscaled):
         if unscaled:
             return act
-        return act*DeformEnv.WORKSPACE_BOX_SIZE
+        return act * DeformEnv.WORKSPACE_BOX_SIZE
 
     def load_objects(self, sim, args, debug):
         res = super(DeformRobotEnv, self).load_objects(sim, args, debug)
@@ -78,15 +78,15 @@ def make_anchors(self):
         preset_dynamic_anchor_vertices = get_preset_properties(
             DEFORM_INFO, self.deform_obj, 'deform_anchor_vertices')
         _, mesh = get_mesh_data(self.sim, self.deform_id)
-        assert(preset_dynamic_anchor_vertices is not None)
+        assert (preset_dynamic_anchor_vertices is not None)
         for i in range(self.num_anchors):  # make anchors
             anchor_pos = np.array(mesh[preset_dynamic_anchor_vertices[i][0]])
             if not np.isfinite(anchor_pos).all():
                 print('anchor_pos not sane:', anchor_pos)
                 input('Press enter to exit')
                 exit(1)
-            link_id = self.robot.info.ee_link_id if i==0 else \
-                          self.robot.info.left_ee_link_id
+            link_id = self.robot.info.ee_link_id if i == 0 else \
+                self.robot.info.left_ee_link_id
             self.sim.createSoftBodyAnchor(
                 self.deform_id, preset_dynamic_anchor_vertices[i][0],
                 self.robot.info.robot_id, link_id)
@@ -146,12 +146,12 @@ def get_grip_obs(self):
         grip_obs = []
         ee_pos, _, ee_linvel, _ = self.robot.get_ee_pos_ori_vel()
         grip_obs.extend(ee_pos)
-        grip_obs.extend((np.array(ee_linvel)/DeformEnv.MAX_OBS_VEL))
+        grip_obs.extend((np.array(ee_linvel) / DeformEnv.MAX_OBS_VEL))
         if self.num_anchors > 1:  # EE pos, vel of left arm
             left_ee_pos, _, left_ee_linvel, _ = \
                 self.robot.get_ee_pos_ori_vel(left=True)
             grip_obs.extend(left_ee_pos)
-            grip_obs.extend((np.array(left_ee_linvel)/DeformEnv.MAX_OBS_VEL))
+            grip_obs.extend((np.array(left_ee_linvel) / DeformEnv.MAX_OBS_VEL))
 
         return grip_obs
 

dedo/utils/args.py

@@ -16,7 +16,6 @@
 import re
 
 
-
 def get_args_parser():
     parser = argparse.ArgumentParser(description='args', add_help=True)
     #
@@ -133,7 +132,7 @@ def get_args_parser():
                         help='Episode rollout length')
     parser.add_argument('--replay_size', type=int, default=10000,
                         help='Number of observations to store in replay buffer'
-                        '10K 200x200 frames take ~20GBs of CPU RAM')
+                             '10K 200x200 frames take ~20GBs of CPU RAM')
     parser.add_argument('--unsup_algo', type=str, default=None,
                         choices=['VAE', 'SVAE', 'PRED', 'DSA'],
                         help='Unsupervised learner (e.g. for run_svae.py)')
@@ -168,6 +167,7 @@ def get_args():
     args_postprocess(args)
     return args
 
+
 def preset_override_util(args, deform_obj):
     '''
     Overrides args object with preset information (deform_obj).
@@ -179,9 +179,9 @@ def preset_override_util(args, deform_obj):
     for argv in sys.argv:
         m = re.search("(?:--)([a-zA-Z0-9-_]+)(?:=)?", argv)
         if m is not None:
-            user_raw_args.append(m.group(1)) # gets the var name
+            user_raw_args.append(m.group(1))  # gets the var name
 
     for arg_nm, arg_val in deform_obj.items():
-        if arg_nm in user_raw_args:   # User overrides the preset arg
+        if arg_nm in user_raw_args:  # User overrides the preset arg
             continue
         setattr(args, arg_nm, arg_val)

dedo/utils/procedural_utils.py

@@ -166,7 +166,6 @@ def gen_random_hole(node_density, dim_constraints):
 
 
 def try_gen_holes(node_density, num_holes, constraints):
-
     for i in range(1000):  # 1000 MC
         if num_holes == 2:
             holeA = gen_random_hole(node_density, constraints)

dedo/utils/rl_sb3_utils.py

@@ -39,7 +39,7 @@ def play(env, num_episodes, rl_agent, debug=False, logdir=None,
         if vidwriter is not None:
             img = env.render(mode='rgb_array', width=cam_resolution,
                              height=cam_resolution)
-            vidwriter.write(img[...,::-1])
+            vidwriter.write(img[..., ::-1])
         step = 0
         while True:
             # rl_agent.predict() to get acts, not forcing deterministic.
@@ -49,7 +49,7 @@ def play(env, num_episodes, rl_agent, debug=False, logdir=None,
             if vidwriter is not None:
                 img = env.render(mode='rgb_array', width=cam_resolution,
                                  height=cam_resolution)
-                vidwriter.write(img[...,::-1])
+                vidwriter.write(img[..., ::-1])
             if done:
                 if debug:
                     print(f'episode reward {episode_rwd:0.2f}')
@@ -65,6 +65,7 @@ class CustomCallback(BaseCallback):
     """
     A custom callback that runs eval and adds videos to Tensorboard.
     """
+
     def __init__(self, eval_env, logdir, num_train_envs, args,
                  num_steps_between_save=10000, viz=False, debug=False):
         super(CustomCallback, self).__init__(debug)
@@ -139,7 +140,7 @@ def grab_screens(_locals, _globals):
 
                 evaluate_policy(
                     self.model, self._eval_env, callback=grab_screens,
-                    n_eval_episodes=1,  deterministic=False)
+                    n_eval_episodes=1, deterministic=False)
                 self.logger.record(
                     'trajectory/video',
                     Video(torch.ByteTensor([screens]), fps=50),
@@ -158,4 +159,3 @@ def _on_training_end(self) -> None:
         This event is triggered before exiting the `learn()` method.
         """
         pass
-

dedo/utils/task_info.py

@@ -224,7 +224,7 @@
                 'mass': 0,
                 'useTexture': True,
             },
-            #'ycb/004_sugar_box/google_16k/textured.obj': {
+            # 'ycb/004_sugar_box/google_16k/textured.obj': {
             # 'ycb/009_gelatin_box/google_16k/textured.obj': {
             'ycb/003_cracker_box/google_16k/textured.obj': {
                 'basePosition': [1.8, 1.7, 0.25],
@@ -233,7 +233,7 @@
                 'mass': 0.01,
                 'rgbaColor': (0.9, 0.75, 0.65, 1),
             },
-            #'ycb/005_tomato_soup_can/google_16k/textured.obj': {
+            # 'ycb/005_tomato_soup_can/google_16k/textured.obj': {
             # 'ycb/007_tuna_fish_can/google_16k/textured.obj': {
             'ycb/002_master_chef_can/google_16k/textured.obj': {
                 'basePosition': [0.9, 1.5, 0.25],
@@ -1190,7 +1190,6 @@
             [0],
             [38],
 
-
         ],
         'deform_fixed_anchor_vertex_ids':
             [372, 373, 374, 375, 376, 377, 378, 379, 380, 381, 382, 383, 384,
@@ -1228,10 +1227,10 @@
     },
     'berkeley_garments': {
         'deform_init_pos': [0, 5, 8],
-        'deform_init_ori': [np.pi/2, 0, 0],
+        'deform_init_ori': [np.pi / 2, 0, 0],
         'deform_scale': 5,
         'anchor_init_pos': [1, 4.4722, 10.4271],
-        'other_anchor_init_pos': [-1,  4.4708, 10.4309],
+        'other_anchor_init_pos': [-1, 4.4708, 10.4309],
         'deform_elastic_stiffness': 10.0,
         'deform_bending_stiffness': 0.05,
         'deform_damping_stiffness': 0.01,
@@ -1242,10 +1241,10 @@
     },
     'sewing_garments': {
         'deform_init_pos': [0, 5, 8],
-        'deform_init_ori': [np.pi/2, 0, 0],
+        'deform_init_ori': [np.pi / 2, 0, 0],
         'deform_scale': 0.05,
         'anchor_init_pos': [1, 4.4722, 10.4271],
-        'other_anchor_init_pos': [-1,  4.4708, 10.4309],
+        'other_anchor_init_pos': [-1, 4.4708, 10.4309],
         'deform_elastic_stiffness': 10.0,
         'deform_bending_stiffness': 0.05,
         'deform_damping_stiffness': 0.01,
@@ -1320,7 +1319,6 @@
                          0.0, 0.0, -0.02000020071864128, 0.0)
 }
 
-
 ROBOT_INFO = {
     'franka2': {
         'file_name': 'franka/franka_dual.urdf',
@@ -1332,10 +1330,10 @@
         'use_fixed_base': True,
         'base_pos': np.array([5.0, 1.5, 0]),
         'rest_arm_qpos': np.array(
-           [-0.7332, -0.0135,  0.1112, -0.718,   0.0978, 1.99, -0.5592]),
+            [-0.7332, -0.0135, 0.1112, -0.718, 0.0978, 1.99, -0.5592]),
         'left_rest_arm_qpos': np.array(
-            [0.7732, -0.0135,  -0.0212, -0.68,   0.0978, 1.99, -0.5592]),
-   },
+            [0.7732, -0.0135, -0.0212, -0.68, 0.0978, 1.99, -0.5592]),
+    },
     'franka1': {
         'file_name': 'franka/franka_small_fingers.urdf',
         'ee_joint_name': 'panda_joint7',
@@ -1344,8 +1342,8 @@
         'use_fixed_base': True,
         'base_pos': np.array([5.0, 1.5, 0]),
         'rest_arm_qpos': np.array(
-           # for [2.5, 1.5, 1.0]
-           [0.4083,  0.4691, -0.6216, -2.9606, -0.9926,  3.4903,  1.5129]
+            # for [2.5, 1.5, 1.0]
+            [0.4083, 0.4691, -0.6216, -2.9606, -0.9926, 3.4903, 1.5129]
         ),
-   }
+    }
 }