import time
import matplotlib.pyplot as plt
from environments import ReacherRandomizedEnv


def test_env(env, n_rolls=10, n_steps=50):
    for i in range(n_rolls):
        env.reset()
        for t in range(n_steps):
            env.step(env.action_space.sample())
            env.render(width=100, height=100)
    # env.render(close=True)
    env.close()


env = ReacherRandomizedEnv()
# start = time.time()
# test_env(env)
# end = time.time()
# print("Total Time: ", end - start)
for i in range(50):
    env.reset()
    img = env.render("rgb_array")
    plt.imsave("reacher_obstacle_{}.jpg".format(i), img)



# from environments.dnc_environments import create_stochastic
#
# env = create_stochastic("pick")
#
# test_env(env)
