119 lines
3.6 KiB
Python
119 lines
3.6 KiB
Python
import jax.numpy as jnp
|
|
from jax import Array
|
|
from .rl_jit import RLEnv, norm_obs
|
|
|
|
|
|
class MujocoEnv(RLEnv):
|
|
def __init__(
|
|
self, env_name: str = "SwimmerSwimmer6", *args, **kwargs
|
|
):
|
|
from mujoco_playground import registry
|
|
super().__init__(*args, **kwargs)
|
|
self.env_name = env_name
|
|
self.env = registry.load(env_name=env_name)
|
|
|
|
def env_step(self, randkey, env_state, action):
|
|
state = self.env.step(env_state, action)
|
|
obs = state.obs
|
|
if not isinstance(obs, Array):
|
|
if "state" in obs:
|
|
obs = obs["state"]
|
|
else:
|
|
raise ImportError(
|
|
f"This Pytree observation space is not supported yet: {obs}"
|
|
)
|
|
return obs, state, state.reward, state.done.astype(jnp.bool_), state.info
|
|
|
|
def env_reset(self, randkey):
|
|
init_state = self.env.reset(randkey)
|
|
obs = init_state.obs
|
|
if not isinstance(obs, Array):
|
|
if "state" in obs:
|
|
obs = obs["state"]
|
|
else:
|
|
raise ImportError(
|
|
f"This Pytree observation space is not supported yet: {obs}"
|
|
)
|
|
return obs, init_state
|
|
|
|
@property
|
|
def input_shape(self):
|
|
return (self.env.observation_size,)
|
|
|
|
@property
|
|
def output_shape(self):
|
|
return (self.env.action_size,)
|
|
|
|
def show(
|
|
self,
|
|
state,
|
|
randkey,
|
|
act_func,
|
|
params,
|
|
save_path=None,
|
|
height=480,
|
|
width=480,
|
|
output_type="rgb_array",
|
|
*args,
|
|
**kwargs,
|
|
):
|
|
|
|
assert output_type in ["gif", "mp4"]
|
|
|
|
import jax
|
|
import imageio
|
|
from brax.io import image
|
|
import numpy as np
|
|
|
|
obs, env_state = self.reset(randkey)
|
|
reward, done = 0.0, False
|
|
state_histories = [env_state.pipeline_state]
|
|
|
|
def step(key, env_state, obs):
|
|
key, _ = jax.random.split(key)
|
|
|
|
if self.obs_normalization:
|
|
obs = norm_obs(state, obs)
|
|
|
|
if self.action_policy is not None:
|
|
forward_func = lambda obs: act_func(state, params, obs)
|
|
action = self.action_policy(key, forward_func, obs)
|
|
else:
|
|
action = act_func(state, params, obs)
|
|
|
|
obs, env_state, r, done, info = self.step(randkey, env_state, action)
|
|
return key, env_state, obs, r, done
|
|
|
|
jit_step = jax.jit(step)
|
|
|
|
for _ in range(self.max_step):
|
|
key, env_state, obs, r, done = jit_step(randkey, env_state, obs)
|
|
state_histories.append(env_state.pipeline_state)
|
|
reward += r
|
|
if done:
|
|
break
|
|
|
|
print("Total reward: ", reward)
|
|
|
|
try:
|
|
imgs = image.render_array(
|
|
sys=self.env.sys, trajectory=state_histories, height=height, width=width, camera="track"
|
|
)
|
|
except ValueError:
|
|
imgs = image.render_array(
|
|
sys=self.env.sys, trajectory=state_histories, height=height, width=width
|
|
)
|
|
|
|
if save_path is None:
|
|
save_path = f"{self.env_name}.{output_type}"
|
|
|
|
imageio.mimsave(save_path, imgs, *args, **kwargs)
|
|
|
|
if output_type == "gif":
|
|
imageio.mimsave(save_path, imgs, *args, **kwargs)
|
|
elif output_type == "mp4":
|
|
fps = kwargs.get("fps", 30)
|
|
imageio.mimsave(save_path, imgs, fps=fps, codec="libx264", format="mp4")
|
|
|
|
print(f"{output_type} saved to: ", save_path)
|