LeroyWaa commited on
Commit
e0cb411
·
1 Parent(s): aa208b9

no robomimic

Browse files
Files changed (1) hide show
  1. 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
- class RobomimicSimulator(PhysicsSimulator):
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