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()
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()
def test_vectorized_original_equality(self):
venv = VectorizedPendulumEnv()
state, action = self.state_action
action = np.round(action)
dim1, dim2 = self.dims
venv.state = state
vobs, vreward, vdone, _ = venv.step(action)
env = PendulumEnv()
for i in range(dim1):
for j in range(dim2):
env.state = state[i, j]
obs, reward, done, _ = env.step(action[i, j])
np.testing.assert_allclose(obs, vobs[i, j])
np.testing.assert_allclose(reward, vreward[i, j])
np.testing.assert_allclose(done, vdone[i, j])
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()