def viz_pendulum_rollout(states, actions):
assert states.shape[0] == actions.shape[0]
eps = jp.finfo(float).eps
gymenv = PendulumEnv()
gymenv.reset()
for t in range(states.shape[0]):
gymenv.state = states[t] + jp.pi
# array(0.0) is False-y which causes problems.
gymenv.last_u = actions[t] + eps
gymenv.render()
gymenv.close()
class GYMMB_Pendulum(Env):
metadata = PendulumEnv.metadata
def __init__(self):
self.env = PendulumEnv()
self.action_space = self.env.action_space
self.observation_space = self.env.observation_space
def reset(self):
return self.env.reset()
def step(self, action):
ob, _, _, info = self.env.step(action)
return ob, None, None, info
def seed(self, seed=None):
return self.env.seed(seed)
def render(self, mode='human'):
return self.env.render(mode)
def close(self):
return self.env.close()
@staticmethod
def tasks():
return dict(standard=StandardTask(), poplin=StandardTask())
@staticmethod
def is_done(states):
bs = states.shape[0]
return torch.zeros(size=[bs], dtype=torch.bool, device=states.device) # Always False
def record_pendulum_rollout(filepath, states, actions):
assert states.shape[0] == actions.shape[0]
eps = jnp.finfo(float).eps
gymenv = PendulumEnv()
gymenv.reset()
video = VideoRecorder(gymenv, path=filepath)
for t in range(states.shape[0]):
gymenv.state = states[t] + jnp.pi
# array(0.0) is False-y which causes problems.
gymenv.last_u = actions[t] + eps
# gymenv.step()
gymenv.render()
video.capture_frame()
video.close()
import numpy as np
from gym.envs.classic_control import PendulumEnv
from gym.wrappers.monitoring.video_recorder import VideoRecorder
if __name__ == "__main__":
T = 1000
gymenv = PendulumEnv()
gymenv.reset()
# Force the initialization.
gymenv.state = [np.pi, 0]
gymenv.last_u = None
video = VideoRecorder(gymenv, path="openai_gym_constant_torque.mp4")
for t in range(T):
gymenv.step([-2])
gymenv.render()
video.capture_frame()
gymenv.close()
video.close()