wrap dm_control aloha into gymnasium (TODO: properly seeding the env)

2024-04-05 17:17:31 +00:00 · 2024-04-05 17:17:31 +00:00 · a2d3588fca
parent 29032fbcd3
commit a2d3588fca
1 changed files with 15 additions and 0 deletions
--- a/lerobot/common/envs/aloha/init.py
+++ b/lerobot/common/envs/aloha/init.py
@ -0,0 +1,15 @@
+from gymnasium.envs.registration import register
+
+register(
+    id="gym_aloha/AlohaInsertion-v0",
+    entry_point="lerobot.common.envs.aloha.env:AlohaEnv",
+    max_episode_steps=300,
+    kwargs={"obs_type": "state", "task": "insertion"},
+)
+
+register(
+    id="gym_aloha/AlohaTransferCube-v0",
+    entry_point="lerobot.common.envs.aloha.env:AlohaEnv",
+    max_episode_steps=300,
+    kwargs={"obs_type": "state", "task": "transfer_cube"},
+)