dodox/dppo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

squash commits

Browse Source
This commit is contained in:
allenzren
2024-09-11 21:09:17 -04:00
parent 8ce0aa1485
commit 2ddf63b8f5
200 changed files with 1240 additions and 1186 deletions
Download Patch File Download Diff File Expand all files Collapse all files
env/gym_utils/wrapper/d3il_lowdim.py
Vendored
+14 -16
View File
@@ -1,10 +1,12 @@
"""
Environment wrapper for D3IL environments with state observations.
For consistency, we will use Dict{} for the observation space, with the key "state" for the state observation.
"""
import numpy as np
import gym
from gym import spaces
class D3ilLowdimWrapper(gym.Env):
@@ -12,28 +14,27 @@ class D3ilLowdimWrapper(gym.Env):
self,
env,
normalization_path,
# init_state=None,
# render_hw=(256, 256),
# render_camera_name="agentview",
):
self.env = env
# self.init_state = init_state
# self.render_hw = render_hw
# self.render_camera_name = render_camera_name
# setup spaces
self.action_space = env.action_space
self.observation_space = env.observation_space
normalization = np.load(normalization_path)
self.obs_min = normalization["obs_min"]
self.obs_max = normalization["obs_max"]
self.action_min = normalization["action_min"]
self.action_max = normalization["action_max"]
# def get_observation(self):
# raw_obs = self.env.get_observation()
# obs = np.concatenate([raw_obs[key] for key in self.obs_keys], axis=0)
# return obs
self.observation_space = spaces.Dict()
obs_example = self.env.reset()
low = np.full_like(obs_example, fill_value=-1)
high = np.full_like(obs_example, fill_value=1)
self.observation_space["state"] = spaces.Box(
low=low,
high=high,
shape=low.shape,
dtype=low.dtype,
)
def seed(self, seed=None):
if seed is not None:
@@ -48,9 +49,6 @@ class D3ilLowdimWrapper(gym.Env):
new_seed = options.get(
"seed", None
) # used to set all environments to specified seeds
# if self.init_state is not None:
# # always reset to the same state to be compatible with gym
# self.env.reset_to({"states": self.init_state})
if new_seed is not None:
self.seed(seed=new_seed)
obs = self.env.reset()
@@ -60,7 +58,7 @@ class D3ilLowdimWrapper(gym.Env):
# normalize
obs = self.normalize_obs(obs)
return obs
return {"state": obs}
def normalize_obs(self, obs):
return 2 * ((obs - self.obs_min) / (self.obs_max - self.obs_min + 1e-6) - 0.5)
@@ -75,7 +73,7 @@ class D3ilLowdimWrapper(gym.Env):
# normalize
obs = self.normalize_obs(obs)
return obs, reward, done, info
return {"state": obs}, reward, done, info
def render(self, mode="rgb_array"):
h, w = self.render_hw
env/gym_utils/wrapper/furniture.py
Vendored
+2 -4
View File
@@ -71,8 +71,7 @@ class FurnitureRLSimEnvMultiStepWrapper(gym.Wrapper):
nobs = self.process_obs(obs)
self.best_reward = torch.zeros(self.env.num_envs).to(self.device)
self.done = list()
return nobs
return {"state": nobs}
def reset_arg(self, options_list=None):
return self.reset()
@@ -80,7 +79,6 @@ class FurnitureRLSimEnvMultiStepWrapper(gym.Wrapper):
def reset_one_arg(self, env_ind=None, options=None):
if env_ind is not None:
env_ind = torch.tensor([env_ind], device=self.device)
return self.reset()
def step(self, action: np.ndarray):
@@ -109,7 +107,7 @@ class FurnitureRLSimEnvMultiStepWrapper(gym.Wrapper):
nobs: np.ndarray = self.process_obs(obs)
done: np.ndarray = done.squeeze().cpu().numpy()
return (nobs, reward, done, info)
return {"state": nobs}, reward, done, info
def _inner_step(self, action_chunk: torch.Tensor):
dones = torch.zeros(
env/gym_utils/wrapper/mujoco_locomotion_lowdim.py
Vendored
+15 -3
View File
@@ -1,10 +1,12 @@
"""
Environment wrapper for Gym environments (MuJoCo locomotion tasks) with state observations.
For consistency, we will use Dict{} for the observation space, with the key "state" for the state observation.
"""
import numpy as np
import gym
from gym import spaces
class MujocoLocomotionLowdimWrapper(gym.Env):
@@ -17,13 +19,23 @@ class MujocoLocomotionLowdimWrapper(gym.Env):
# setup spaces
self.action_space = env.action_space
self.observation_space = env.observation_space
normalization = np.load(normalization_path)
self.obs_min = normalization["obs_min"]
self.obs_max = normalization["obs_max"]
self.action_min = normalization["action_min"]
self.action_max = normalization["action_max"]
self.observation_space = spaces.Dict()
obs_example = self.env.reset()
low = np.full_like(obs_example, fill_value=-1)
high = np.full_like(obs_example, fill_value=1)
self.observation_space["state"] = spaces.Box(
low=low,
high=high,
shape=low.shape,
dtype=low.dtype,
)
def seed(self, seed=None):
if seed is not None:
np.random.seed(seed=seed)
@@ -40,7 +52,7 @@ class MujocoLocomotionLowdimWrapper(gym.Env):
# normalize
obs = self.normalize_obs(raw_obs)
return obs
return {"state": obs}
def normalize_obs(self, obs):
return 2 * ((obs - self.obs_min) / (self.obs_max - self.obs_min + 1e-6) - 0.5)
@@ -55,7 +67,7 @@ class MujocoLocomotionLowdimWrapper(gym.Env):
# normalize
obs = self.normalize_obs(raw_obs)
return obs, reward, done, info
return {"state": obs}, reward, done, info
def render(self, **kwargs):
return self.env.render()
env/gym_utils/wrapper/multi_step.py
Vendored
+2 -20
View File
@@ -3,6 +3,7 @@ Multi-step wrapper. Allow executing multiple environmnt steps. Returns stacked o
Modified from https://github.com/real-stanford/diffusion_policy/blob/main/diffusion_policy/gym_util/multistep_wrapper.py
TODO: allow cond_steps != img_cond_steps (should be implemented in training scripts, not here)
"""
import gym
@@ -11,8 +12,6 @@ from gym import spaces
import numpy as np
from collections import defaultdict, deque
# import dill
def stack_repeated(x, n):
return np.repeat(np.expand_dims(x, axis=0), n, axis=0)
@@ -157,12 +156,11 @@ class MultiStep(gym.Wrapper):
done = True
self.done.append(done)
self._add_info(info)
observation = self._get_obs(self.n_obs_steps)
reward = aggregate(self.reward, self.reward_agg_method)
done = aggregate(self.done, "max")
info = dict_take_last_n(self.info, self.n_obs_steps)
if self.pass_full_observations: # right now this assume n_obs_steps = 1
if self.pass_full_observations:
info["full_obs"] = self._get_obs(act_step + 1)
# In mujoco case, done can happen within the loop above
@@ -206,22 +204,6 @@ class MultiStep(gym.Wrapper):
"""Not the best design"""
return self.env.render(**kwargs)
# def get_rewards(self):
# return self.reward
# def get_attr(self, name):
# return getattr(self, name)
# def run_dill_function(self, dill_fn):
# fn = dill.loads(dill_fn)
# return fn(self)
# def get_infos(self):
# result = dict()
# for k, v in self.info.items():
# result[k] = list(v)
# return result
if __name__ == "__main__":
import os
env/gym_utils/wrapper/robomimic_image.py
Vendored
+1 -3
View File
@@ -90,9 +90,7 @@ class RobomimicImageWrapper(gym.Env):
action = (action + 1) / 2 # [-1, 1] -> [0, 1]
return action * (self.action_max - self.action_min) + self.action_min
def get_observation(self, raw_obs=None):
if raw_obs is None:
raw_obs = self.env.get_observation()
def get_observation(self, raw_obs):
obs = {"rgb": None, "state": None} # stack rgb if multiple cameras
for key in self.obs_keys:
if key in self.image_keys:
env/gym_utils/wrapper/robomimic_lowdim.py
Vendored
+20 -21
View File
@@ -3,11 +3,12 @@ Environment wrapper for Robomimic environments with state observations.
Modified from https://github.com/real-stanford/diffusion_policy/blob/main/diffusion_policy/env/robomimic/robomimic_lowdim_wrapper.py
For consistency, we will use Dict{} for the observation space, with the key "state" for the state observation.
"""
import numpy as np
import gym
from gym.spaces import Box
from gym import spaces
import imageio
@@ -28,7 +29,6 @@ class RobomimicLowdimWrapper(gym.Env):
render_camera_name="agentview",
):
self.env = env
self.obs_keys = low_dim_keys
self.init_state = init_state
self.render_hw = render_hw
self.render_camera_name = render_camera_name
@@ -44,19 +44,24 @@ class RobomimicLowdimWrapper(gym.Env):
self.action_min = normalization["action_min"]
self.action_max = normalization["action_max"]
# setup spaces - use [-1, 1]
# setup spaces
low = np.full(env.action_dimension, fill_value=-1)
high = np.full(env.action_dimension, fill_value=1)
self.action_space = Box(
self.action_space = gym.spaces.Box(
low=low,
high=high,
shape=low.shape,
dtype=low.dtype,
)
obs_example = self.get_observation()
self.obs_keys = low_dim_keys
self.observation_space = spaces.Dict()
obs_example_full = self.env.get_observation()
obs_example = np.concatenate(
[obs_example_full[key] for key in self.obs_keys], axis=0
)
low = np.full_like(obs_example, fill_value=-1)
high = np.full_like(obs_example, fill_value=1)
self.observation_space = Box(
self.observation_space["state"] = spaces.Box(
low=low,
high=high,
shape=low.shape,
@@ -75,12 +80,11 @@ class RobomimicLowdimWrapper(gym.Env):
action = (action + 1) / 2 # [-1, 1] -> [0, 1]
return action * (self.action_max - self.action_min) + self.action_min
def get_observation(self):
raw_obs = self.env.get_observation()
raw_obs = np.concatenate([raw_obs[key] for key in self.obs_keys], axis=0)
def get_observation(self, raw_obs):
obs = {"state": np.concatenate([raw_obs[key] for key in self.obs_keys], axis=0)}
if self.normalize:
return self.normalize_obs(raw_obs)
return raw_obs
obs["state"] = self.normalize_obs(obs["state"])
return obs
def seed(self, seed=None):
if seed is not None:
@@ -90,7 +94,6 @@ class RobomimicLowdimWrapper(gym.Env):
def reset(self, options={}, **kwargs):
"""Ignore passed-in arguments like seed"""
# Close video if exists
if self.video_writer is not None:
self.video_writer.close()
@@ -106,24 +109,20 @@ class RobomimicLowdimWrapper(gym.Env):
) # used to set all environments to specified seeds
if self.init_state is not None:
# always reset to the same state to be compatible with gym
self.env.reset_to({"states": self.init_state})
raw_obs = self.env.reset_to({"states": self.init_state})
elif new_seed is not None:
self.seed(seed=new_seed)
self.env.reset()
raw_obs = self.env.reset()
else:
# random reset
self.env.reset()
return self.get_observation()
raw_obs = self.env.reset()
return self.get_observation(raw_obs)
def step(self, action):
if self.normalize:
action = self.unnormalize_action(action)
raw_obs, reward, done, info = self.env.step(action)
raw_obs = np.concatenate([raw_obs[key] for key in self.obs_keys], axis=0)
if self.normalize:
obs = self.normalize_obs(raw_obs)
else:
obs = raw_obs
obs = self.get_observation(raw_obs)
# render if specified
if self.video_writer is not None: