Spaces:
Running
on
Zero
Running
on
Zero
no robomimic
Browse files- sim/simulator.py +1 -62
sim/simulator.py
CHANGED
@@ -4,10 +4,7 @@ import numpy as np
|
|
4 |
import einops
|
5 |
import skimage
|
6 |
import time
|
7 |
-
import robomimic.utils.file_utils as FileUtils
|
8 |
|
9 |
-
from sim.robomimic.robomimic_wrapper import RobomimicLowdimWrapper
|
10 |
-
from sim.robomimic.robomimic_runner import create_env, OBS_KEYS, RESOLUTION
|
11 |
from genie.st_mask_git import STMaskGIT
|
12 |
from genie.st_mar import STMAR
|
13 |
from datasets.utils import get_image_encoder
|
@@ -103,65 +100,7 @@ class ReplaySimulator(PhysicsSimulator):
|
|
103 |
return self._dt
|
104 |
|
105 |
|
106 |
-
|
107 |
-
"""
|
108 |
-
Example usage:
|
109 |
-
sim = RobomimicSimulator("lift")
|
110 |
-
useless_low_dim_state = sim.reset()
|
111 |
-
image = sim.step([0] * 7)
|
112 |
-
Image.fromarray(image).save("test.png")
|
113 |
-
"""
|
114 |
-
def __init__(self, env_name):
|
115 |
-
super().__init__()
|
116 |
-
env_meta = FileUtils.get_env_metadata_from_dataset(
|
117 |
-
f'data/robomimic_datasets/robomimic_raw/datasets/{env_name}/ph/image.hdf5')
|
118 |
-
env_meta['use_image_obs'] = True
|
119 |
-
env = create_env(env_meta, obs_keys=OBS_KEYS)
|
120 |
-
self.env = RobomimicLowdimWrapper(env=env, obs_keys=OBS_KEYS)
|
121 |
-
self.env_name = env_name
|
122 |
-
|
123 |
-
def set_initial_state(self, state):
|
124 |
-
self.env.init_state = state
|
125 |
-
|
126 |
-
def step(self, action):
|
127 |
-
obs, reward, done, info = self.env.step(action)
|
128 |
-
img = self.env.render()
|
129 |
-
img = cv2.resize(img, RESOLUTION)
|
130 |
-
return {
|
131 |
-
'pred_next_frame': img,
|
132 |
-
'reward': reward,
|
133 |
-
'done': done,
|
134 |
-
'info': info
|
135 |
-
}
|
136 |
-
|
137 |
-
def reset(self):
|
138 |
-
self.env.reset() # reset to initial state
|
139 |
-
img = self.env.render()
|
140 |
-
img = cv2.resize(img, RESOLUTION)
|
141 |
-
return img
|
142 |
-
|
143 |
-
def close(self):
|
144 |
-
self.env.close()
|
145 |
-
|
146 |
-
@property
|
147 |
-
def dt(self):
|
148 |
-
return self.env.env.env.sim.model.opt.timestep
|
149 |
-
|
150 |
-
@property
|
151 |
-
def action_dimension(self):
|
152 |
-
return self.env.env.action_dimension
|
153 |
-
|
154 |
-
def set_dt(self, dt):
|
155 |
-
self.env.env.env.sim.model.opt.timestep = dt
|
156 |
-
|
157 |
-
def get_raw_state(self, port: Optional[str] = None):
|
158 |
-
raw_state = self.env.env.get_state()
|
159 |
-
raw_obs = self.env.env.get_observation()
|
160 |
-
raw_state.update(raw_obs)
|
161 |
-
if port is not None:
|
162 |
-
return raw_state[port]
|
163 |
-
return raw_state
|
164 |
-
|
165 |
|
166 |
class GenieSimulator(LearnedSimulator):
|
167 |
|
|
|
4 |
import einops
|
5 |
import skimage
|
6 |
import time
|
|
|
7 |
|
|
|
|
|
8 |
from genie.st_mask_git import STMaskGIT
|
9 |
from genie.st_mar import STMAR
|
10 |
from datasets.utils import get_image_encoder
|
|
|
100 |
return self._dt
|
101 |
|
102 |
|
103 |
+
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
104 |
|
105 |
class GenieSimulator(LearnedSimulator):
|
106 |
|