improved project structure and exposed methods

2021-08-25 17:16:20 +02:00
parent 031dba541a
commit c443b06fef
160 changed files with 1010 additions and 1575 deletions
@@ -1 +1,114 @@
-from alr_envs.meta import goal_and_object_change, goal_and_endeffector_change, goal_change, object_change
+from gym import register
+
+from . import goal_object_change_mp_wrapper, goal_change_mp_wrapper, goal_endeffector_change_mp_wrapper, \
+    object_change_mp_wrapper
+
+ALL_METAWORLD_MOTION_PRIMITIVE_ENVIRONMENTS = {"DMP": [], "DetPMP": []}
+
+# MetaWorld
+
+_goal_change_envs = ["assembly-v2", "pick-out-of-hole-v2", "plate-slide-v2", "plate-slide-back-v2"]
+for _task in _goal_change_envs:
+    task_id_split = _task.split("-")
+    name = "".join([s.capitalize() for s in task_id_split[:-1]])
+    _env_id = f'{name}DetPMP-{task_id_split[-1]}'
+    register(
+        id=_env_id,
+        entry_point='alr_envs.utils.make_env_helpers:make_detpmp_env_helper',
+        kwargs={
+            "name": _task,
+            "wrappers": [goal_change_mp_wrapper.MPWrapper],
+            "mp_kwargs": {
+                "num_dof": 4,
+                "num_basis": 5,
+                "duration": 6.25,
+                "post_traj_time": 0,
+                "width": 0.025,
+                "zero_start": True,
+                "policy_type": "metaworld",
+            }
+        }
+    )
+    ALL_METAWORLD_MOTION_PRIMITIVE_ENVIRONMENTS["DetPMP"].append(_env_id)
+
+_object_change_envs = ["bin-picking-v2", "hammer-v2", "sweep-into-v2"]
+for _task in _object_change_envs:
+    task_id_split = _task.split("-")
+    name = "".join([s.capitalize() for s in task_id_split[:-1]])
+    _env_id = f'{name}DetPMP-{task_id_split[-1]}'
+    register(
+        id=_env_id,
+        entry_point='alr_envs.utils.make_env_helpers:make_detpmp_env_helper',
+        kwargs={
+            "name": _task,
+            "wrappers": [object_change_mp_wrapper.MPWrapper],
+            "mp_kwargs": {
+                "num_dof": 4,
+                "num_basis": 5,
+                "duration": 6.25,
+                "post_traj_time": 0,
+                "width": 0.025,
+                "zero_start": True,
+                "policy_type": "metaworld",
+            }
+        }
+    )
+    ALL_METAWORLD_MOTION_PRIMITIVE_ENVIRONMENTS["DetPMP"].append(_env_id)
+
+_goal_and_object_change_envs = ["box-close-v2", "button-press-v2", "button-press-wall-v2", "button-press-topdown-v2",
+                                "button-press-topdown-wall-v2", "coffee-button-v2", "coffee-pull-v2",
+                                "coffee-push-v2", "dial-turn-v2", "disassemble-v2", "door-close-v2",
+                                "door-lock-v2", "door-open-v2", "door-unlock-v2", "hand-insert-v2",
+                                "drawer-close-v2", "drawer-open-v2", "faucet-open-v2", "faucet-close-v2",
+                                "handle-press-side-v2", "handle-press-v2", "handle-pull-side-v2",
+                                "handle-pull-v2", "lever-pull-v2", "peg-insert-side-v2", "pick-place-wall-v2",
+                                "reach-v2", "push-back-v2", "push-v2", "pick-place-v2", "peg-unplug-side-v2",
+                                "soccer-v2", "stick-push-v2", "stick-pull-v2", "push-wall-v2", "reach-wall-v2",
+                                "shelf-place-v2", "sweep-v2", "window-open-v2", "window-close-v2"
+                                ]
+for _task in _goal_and_object_change_envs:
+    task_id_split = _task.split("-")
+    name = "".join([s.capitalize() for s in task_id_split[:-1]])
+    _env_id = f'{name}DetPMP-{task_id_split[-1]}'
+    register(
+        id=_env_id,
+        entry_point='alr_envs.utils.make_env_helpers:make_detpmp_env_helper',
+        kwargs={
+            "name": _task,
+            "wrappers": [goal_object_change_mp_wrapper.MPWrapper],
+            "mp_kwargs": {
+                "num_dof": 4,
+                "num_basis": 5,
+                "duration": 6.25,
+                "post_traj_time": 0,
+                "width": 0.025,
+                "zero_start": True,
+                "policy_type": "metaworld",
+            }
+        }
+    )
+    ALL_METAWORLD_MOTION_PRIMITIVE_ENVIRONMENTS["DetPMP"].append(_env_id)
+
+_goal_and_endeffector_change_envs = ["basketball-v2"]
+for _task in _goal_and_endeffector_change_envs:
+    task_id_split = _task.split("-")
+    name = "".join([s.capitalize() for s in task_id_split[:-1]])
+    _env_id = f'{name}DetPMP-{task_id_split[-1]}'
+    register(
+        id=_env_id,
+        entry_point='alr_envs.utils.make_env_helpers:make_detpmp_env_helper',
+        kwargs={
+            "name": _task,
+            "wrappers": [goal_endeffector_change_mp_wrapper.MPWrapper],
+            "mp_kwargs": {
+                "num_dof": 4,
+                "num_basis": 5,
+                "duration": 6.25,
+                "post_traj_time": 0,
+                "width": 0.025,
+                "zero_start": True,
+                "policy_type": "metaworld",
+            }
+        }
+    )
+    ALL_METAWORLD_MOTION_PRIMITIVE_ENVIRONMENTS["DetPMP"].append(_env_id)