update omni worker

StanfordVL · Nov 22, 2023 · d0a17d6 · d0a17d6
1 parent a7ae750
commit d0a17d6
Showing 1 changed file with 14 additions and 49 deletions.
diff --git a/rl/service/omni_grpc_worker.py b/rl/service/omni_grpc_worker.py
@@ -1,16 +1,16 @@
 import asyncio
+import numpy as np
 import omnigibson as og
-from omnigibson.envs.rl_env import RLEnv
 from omnigibson.macros import gm
-import h5py
 
 from rollout_worker import serve
 
+gm.USE_GPU_DYNAMICS = True
 
 async def main():
+
     DIST_COEFF = 0.1
     GRASP_REWARD = 0.3
-    h5py.get_config().track_order = True
 
     cfg = {
         "scene": {
@@ -20,9 +20,9 @@ async def main():
         },
         "robots": [
             {
-                "type": "Tiago",
+                "type": "Fetch",
                 "obs_modalities": ["rgb", "depth_linear", "seg_instance", "seg_semantic", "proprio"],
-                "proprio_obs": ["robot_pose", "joint_qpos", "joint_qvel", "eef_left_pos", "eef_left_quat", "grasp_left"],
+                "proprio_obs": ["robot_pos", "robot_2d_ori", "joint_qpos", "joint_qvel", "eef_0_pos", "eef_0_quat", "grasp_0"],
                 "scale": 1.0,
                 "self_collisions": True,
                 "action_normalize": False,
@@ -42,46 +42,23 @@ async def main():
                 },
                 "controller_config": {
                     "base": {
-                        "name": "JointController",
-                        "motor_type": "velocity"
+                        "name": "DifferentialDriveController",
                     },
-                    "arm_left": {
+                    "arm_0": {
                         "name": "InverseKinematicsController",
                         "motor_type": "velocity",
-                        "command_input_limits": None,
+                        "command_input_limits": (np.array([-0.2, -0.2, -0.2, -np.pi, -np.pi, -np.pi]),
+                        np.array([0.2, 0.2, 0.2, np.pi, np.pi, np.pi])),
                         "command_output_limits": None,
                         "mode": "pose_absolute_ori", 
                         "kv": 3.0
                     },
-                    # "arm_left": {
-                    #     "name": "JointController",
-                    #     "motor_type": "position",
-                    #     "command_input_limits": None,
-                    #     "command_output_limits": None, 
-                    #     "use_delta_commands": False
-                    # },
-                    "arm_right": {
-                        "name": "JointController",
-                        "motor_type": "position",
-                        "command_input_limits": None,
-                        "command_output_limits": None, 
-                        "use_delta_commands": False
-                    },
-                    "gripper_left": {
-                        "name": "JointController",
+                    "gripper_0": {
+                        "name": "MultiFingerGripperController",
                         "motor_type": "position",
                         "command_input_limits": [-1, 1],
                         "command_output_limits": None,
-                        "use_delta_commands": True,
-                        "use_single_command": True
-                    },
-                    "gripper_right": {
-                        "name": "JointController",
-                        "motor_type": "position",
-                        "command_input_limits": [-1, 1],
-                        "command_output_limits": None,
-                        "use_delta_commands": True,
-                        "use_single_command": True
+                        "use_delta_commands": True
                     },
                     "camera": {
                         "name": "JointController",
@@ -97,7 +74,7 @@ async def main():
             "type": "GraspTask",
             "obj_name": "cologne",
             "termination_config": {
-                "max_steps": 100000,
+                "max_steps": 400,
             },
             "reward_config": {
                 "r_dist_coeff": DIST_COEFF,
@@ -115,19 +92,7 @@ async def main():
         ]
     }
 
-    reset_positions =  {
-        'coffee_table_fqluyq_0': ([-0.4767243 , -1.219805  ,  0.25702515], [-3.69874935e-04, -9.39229270e-04,  7.08872199e-01,  7.05336273e-01]),
-        'cologne': ([-0.30000001, -0.80000001,  0.44277492],
-                    [0.        , 0.        , 0.        , 1.000000]),
-        'robot0': ([0.0, 0.0, 0.05], [0.0, 0.0, 0.0, 1.0])
-    }
-
-    env_config = {
-        "cfg": cfg,
-        "reset_positions": reset_positions,
-        "action_space_controllers": ["base", "camera", "arm_left", "gripper_left"]
-    }
-    env = RLEnv(env_config)
+    env = og.Environment(configs=cfg, action_timestep=1 / 10., physics_timestep=1 / 60.)
 
     # Now start servicing!
     await serve(env)