add visual pointmaze

.gitignore

@@ -6,3 +6,13 @@ dist/
 .DS_Store
 .idea/
 .ruff_cache/
+taming/
+lightning_logs/
+logs/
+wandb/
+data/
+*.pth
+*.gif
+*.pth
+*.npy
+*.npz

data_gen_scripts/generate_locomaze.py

@@ -1,3 +1,17 @@
+import sys
+import os
+sys.path.insert(0, "/home/hyeons/workspace/ogbench") # Path to local ogbench
+sys.path.append("../impls") 
+os.environ["MUJOCO_GL"] = "egl"
+# python generate_locomaze.py --env_name=visual-pointmaze-medium-v0 --save_path=data/visual-pointmaze-stitch-navigate-v0.npz --dataset_type=stitch --num_episodes=5000 --max_episode_steps=201 --noise=0.5
+# python generate_locomaze.py --env_name=visual-pointmaze-giant-v0 --save_path=data/visual-pointmaze-giant-stitch-v0.npz --dataset_type=stitch --num_episodes=5000 --max_episode_steps=201 --noise=0.5
+# python generate_locomaze.py --env_name=visual-pointmaze-medium-v0 --save_path=data/visual-pointmaze-medium-navigate-v0.npz --dataset_type=navigate --num_episodes=1000 --max_episode_steps=1001 --noise=0.5
+# python generate_locomaze.py --env_name=visual-pointmaze-large-v0 --save_path=data/visual-pointmaze-medium-large-v0.npz --dataset_type=navigate --num_episodes=1000 --max_episode_steps=1001 --noise=0.5
+# python generate_locomaze.py --env_name=visual-pointmaze-giant-v0 --save_path=data/visual-pointmaze-giant-large-v0.npz --dataset_type=navigate --num_episodes=500 --max_episode_steps=2001 --noise=0.5
+
+
+# for test
+# python generate_locomaze.py --env_name=visual-pointmaze-giant-v0 --save_path=./data/0331.npz --dataset_type=navigate --num_episodes=1 --max_episode_steps=2001 --noise=0.5
 import glob
 import json
 from collections import defaultdict
@@ -190,7 +204,9 @@ def actor_fn(ob, temperature):
 
     train_path = FLAGS.save_path
     val_path = FLAGS.save_path.replace('.npz', '-val.npz')
-
+    # Check if the directory exists, if not, create it.
+    os.makedirs('/'.join(train_path.split('/')[:-1]), exist_ok=True)
+    print(f'Saving to {train_path} and {val_path}')
     # Split the dataset into training and validation sets.
     train_dataset = {}
     val_dataset = {}

data_gen_scripts/vis.ipynb

@@ -0,0 +1,78 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "import imageio\n",
+    "import IPython.display as display\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 25,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def sample_gif(f_name):\n",
+    "    data = np.load('./data/'+ f_name + '.npz') # load만 26초 정도 걸림\n",
+    "    print('Loaded data from','./data/'+ f_name + '.npz')\n",
+    "    obs = data['observations'][:1000]\n",
+    "    with imageio.get_writer(f_name + '.gif', mode='I') as writer:\n",
+    "        for i in range(0, len(obs), 3):  # Skip every 2 frames to speed up by 3x\n",
+    "            writer.append_data(obs[i])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 26,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "f_name = '0331'\n",
+    "sample_gif(f_name)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import ogbench\n",
+    "env = ogbench.make('visual-pointmaze-giant-stitch-v0')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "og_game",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.15"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

ogbench/locomaze/__init__.py

@@ -19,6 +19,17 @@
         maze_type='medium',
     ),
 )
+register(
+    id='visual-pointmaze-medium-v0',
+    entry_point='ogbench.locomaze.maze:make_maze_env',
+    max_episode_steps=1000,
+    kwargs=dict(
+        loco_env_type='point',
+        maze_env_type='maze',
+        maze_type='medium',
+        **visual_dict,
+    ),
+)
 register(
     id='pointmaze-large-v0',
     entry_point='ogbench.locomaze.maze:make_maze_env',
@@ -29,6 +40,17 @@
         maze_type='large',
     ),
 )
+register(
+    id='visual-pointmaze-large-v0',
+    entry_point='ogbench.locomaze.maze:make_maze_env',
+    max_episode_steps=1000,
+    kwargs=dict(
+        loco_env_type='point',
+        maze_env_type='maze',
+        maze_type='large',
+        **visual_dict,
+    ),
+)
 register(
     id='pointmaze-giant-v0',
     entry_point='ogbench.locomaze.maze:make_maze_env',
@@ -39,6 +61,17 @@
         maze_type='giant',
     ),
 )
+register(
+    id='visual-pointmaze-giant-v0',
+    entry_point='ogbench.locomaze.maze:make_maze_env',
+    max_episode_steps=1000,
+    kwargs=dict(
+        loco_env_type='point',
+        maze_env_type='maze',
+        maze_type='giant',
+        **visual_dict,
+    ),
+)
 register(
     id='pointmaze-teleport-v0',
     entry_point='ogbench.locomaze.maze:make_maze_env',
@@ -49,7 +82,17 @@
         maze_type='teleport',
     ),
 )
-
+register(
+    id='visual-pointmaze-teleport-v0',
+    entry_point='ogbench.locomaze.maze:make_maze_env',
+    max_episode_steps=1000,
+    kwargs=dict(
+        loco_env_type='point',
+        maze_env_type='maze',
+        maze_type='teleport',
+        **visual_dict,
+    ),
+)
 register(
     id='antmaze-medium-v0',
     entry_point='ogbench.locomaze.maze:make_maze_env',

ogbench/locomaze/assets/point.xml

@@ -27,6 +27,10 @@
         <light name="global" directional="true" cutoff="100" exponent="1" ambient=".2 .2 .2" diffuse="1 1 1" specular=".1 .1 .1" pos="0 0 1.3" dir="-0 0 -1.3"/>
         <geom name="floor" pos="0 0 0" size="100 100 .2" type="plane" conaffinity="1" condim="3" material="grid"/>
         <body name="torso" pos="0 0 0">
+            <!-- <camera name="back" pos="0 -2.5 5" xyaxes="1 0 0 0 2 1" mode="trackcom"/> Default -->
+            <!-- <camera name="back" pos="0 -4 6" xyaxes="1 0 0 0 1.5 1" mode="trackcom"/>  on MCTD-->
+            <camera name="back" pos="0 0 15" xyaxes="1 0 0 0 1 0" mode="fixed"/>
+            <camera name="maze-view" pos="0 0 100" xyaxes="1 0 0 0 1 0" mode="fixed"/>
             <geom name="pointbody" type="sphere" size="0.7" pos="0 0 0.7"/>
             <joint name="ballx" type="slide" axis="1 0 0" pos="0 0 0"/>
             <joint name="bally" type="slide" axis="0 1 0" pos="0 0 0"/>

ogbench/locomaze/maze.py

@@ -327,18 +327,18 @@ def set_tasks(self):
                 raise ValueError(f'Unknown maze type: {self._maze_type}')
 
             # More diverse task generation based on the maze map
-            visitable_positions = []
-            for i in range(self.maze_map.shape[0]):
-                for j in range(self.maze_map.shape[1]):
-                    if self.maze_map[i, j] == 0:
-                        visitable_positions.append((i, j))
-            # Combinations
-            tasks = []
-            for i in range(len(visitable_positions)):
-                for j in range(i + 1, len(visitable_positions)):
-                    tasks.append([visitable_positions[i], visitable_positions[j]])
-                    tasks.append([visitable_positions[j], visitable_positions[i]])
-            print(f"The number of tasks is {len(tasks)}. The Task ID should be in [1, {len(tasks)}].")
+            # visitable_positions = []
+            # for i in range(self.maze_map.shape[0]):
+            #     for j in range(self.maze_map.shape[1]):
+            #         if self.maze_map[i, j] == 0:
+            #             visitable_positions.append((i, j))
+            # # Combinations
+            # tasks = []
+            # for i in range(len(visitable_positions)):
+            #     for j in range(i + 1, len(visitable_positions)):
+            #         tasks.append([visitable_positions[i], visitable_positions[j]])
+            #         tasks.append([visitable_positions[j], visitable_positions[i]])
+            # print(f"The number of tasks is {len(tasks)}. The Task ID should be in [1, {len(tasks)}].")
 
             self.task_infos = []
             for i, task in enumerate(tasks):

ogbench/manipspace/envs/manipspace_env.py

@@ -453,5 +453,5 @@ def render(
     ):
         if camera is None:
             camera = 'front' if self._ob_type == 'states' else 'front_pixels'
-
+            
         return super().render(camera=camera, *args, **kwargs)

ogbench/pretrain/dataset/vog_maze.py

@@ -0,0 +1,62 @@
+import torch
+import numpy as np
+
+class VOGMaze2dOfflineRLDataset(torch.utils.data.Dataset):
+    '''
+    Offline RL dataset for 2D maze environments from OG-Bench.
+    Large
+        Mean of obs: 139.58089505083132
+        std of obs: 71.31185013523307
+        Mean of pos: [16.702621 10.974173]
+        std of pos: [10.050303   6.8203936]
+    Giant
+        Mean of obs: 141.01873851323037
+        std of obs: 73.4250522212486
+        Mean of pos: [24.888689 17.158426]
+        std of pos: [14.732276 11.651127]
+    '''
+
+    def __init__(self, dataset_url='/home/hyeons/workspace/ogbench/ogbench/dataset/visual-pointmaze-medium-navigate-v0.npz' , split: str = "training"):
+
+        super().__init__()
+        self.dataset_url = dataset_url
+        self.split = split
+        dataset = self.get_dataset(self.dataset_url)
+        self.observations = dataset["observations"]
+        self.pos = dataset["qpos"]
+        self.actions = dataset["actions"]
+
+        # Normalizations
+        if 'large' in dataset_url:
+            self.observations = (self.observations - 139.58089505083132) / 71.31185013523307
+        elif 'giant' in dataset_url:
+            self.observations = (self.observations - 141.01873851323037) / 73.4250522212486
+
+    def __getitem__(self, idx):
+        observation = torch.from_numpy(self.observations[idx]).float().permute(2, 0, 1)
+        pos = torch.from_numpy(self.pos[idx]).float()
+        action = torch.from_numpy(self.actions[idx]).float()
+        return observation, pos, action
+
+    def __len__(self):
+        return len(self.observations)
+
+    def get_dataset(self, path):
+        if self.split == "validation":
+            path = path.replace(".npz", "-val.npz")
+        dataset = np.load(path, allow_pickle=True, mmap_mode='r')  # 메모리 매핑 적용
+        return dataset
+
+if __name__ == '__main__':
+    dataset = VOGMaze2dOfflineRLDataset(dataset_url='/home/hyeons/workspace/ogbench/data_gen_scripts/data/visual-pointmaze-large-navigate-v0.npz', split='training')
+    print('large Mean of obs:', np.mean(dataset.observations))
+    print('large std of obs:', np.std(dataset.observations))
+    print("large Mean of pos:", np.mean(dataset.pos, axis=0))
+    print("large std of pos:", np.std(dataset.pos, axis=0))
+
+    dataset = VOGMaze2dOfflineRLDataset(dataset_url='/home/hyeons/workspace/ogbench/data_gen_scripts/data/visual-pointmaze-giant-navigate-v0.npz', split='training')
+    print('Giant Mean of obs:', np.mean(dataset.observations))
+    print('Giant std of obs:', np.std(dataset.observations))
+    print("Giant Mean of pos:", np.mean(dataset.pos, axis=0))
+    print("Giant std of pos:", np.std(dataset.pos, axis=0))
+