# Path: osim/env/arm.py
# class Arm2DVecEnv(Arm2DEnv):
#     def reset(self, obs_as_dict=False):
#         obs = super(Arm2DVecEnv, self).reset(obs_as_dict=obs_as_dict)
#         if np.isnan(obs).any():
#             obs = np.nan_to_num(obs)
#         return obs
#     def step(self, action, obs_as_dict=False):
#         if np.isnan(action).any():
#             action = np.nan_to_num(action)
#         obs, reward, done, info = super(Arm2DVecEnv, self).step(action, obs_as_dict=obs_as_dict)
#         if np.isnan(obs).any():
#             obs = np.nan_to_num(obs)
#             done = True
#             reward -10
#         return obs, reward, done, info


import opensim as osim
import numpy as np
import sys
import numpy as np
import argparse
import math
from keras.models import Sequential, Model
from keras.layers import Dense, Activation, Flatten, Input, concatenate
from keras.optimizers import Adam
from rl.agents import DDPGAgent
from rl.memory import SequentialMemory
from rl.random import OrnsteinUhlenbeckProcess
from osim.env.arm import Arm2DVecEnv
from keras.optimizers import RMSprop
# Derived from keras-rl







# Command line parameters
parser = argparse.ArgumentParser(description='Train or test neural net motor controller')
parser.add_argument('--train', dest='train', action='store_true', default=True)
parser.add_argument('--test', dest='train', action='store_false', default=True)
parser.add_argument('--steps', dest='steps', action='store', default=10000, type=int)
parser.add_argument('--visualize', dest='visualize', action='store_true', default=False)
parser.add_argument('--model', dest='model', action='store', default="example.h5f")
args = parser.parse_args()

# set to get observation in array
#def _new_step(self, action, project=True, obs_as_dict=False):
#    return super(Arm2DEnv, self).step(action, project=project, obs_as_dict=obs_as_dict)
#Arm2DEnv.step = _new_step
# Load walking environment

# Path: osim/env/arm.py
# class Arm2DVecEnv(Arm2DEnv):
#     def reset(self, obs_as_dict=False):
#         obs = super(Arm2DVecEnv, self).reset(obs_as_dict=obs_as_dict)
#         if np.isnan(obs).any():
#             obs = np.nan_to_num(obs)
#         return obs
#     def step(self, action, obs_as_dict=False):
#         if np.isnan(action).any():
#             action = np.nan_to_num(action)
#         obs, reward, done, info = super(Arm2DVecEnv, self).step(action, obs_as_dict=obs_as_dict)
#         if np.isnan(obs).any():
#             obs = np.nan_to_num(obs)
#             done = True
#             reward -10
#         return obs, reward, done, info


# Path: examples/train_arm.py
import opensim as osim
import numpy as np
import sys
import numpy as np
import argparse
import math
from keras.models import Sequential, Model
from keras.layers import Dense, Activation, Flatten, Input, concatenate
from keras.optimizers import Adam
from rl.agents import DDPGAgent
from rl.memory import SequentialMemory
from rl.random import OrnsteinUhlenbeckProcess
from osim.env.arm import Arm2DVecEnv
from keras.optimizers import RMSprop

# Derived from keras-rl







# Command line parameters
parser = argparse.ArgumentParser(description='Train or test neural net motor controller')
parser.add_argument('--train', dest='train', action='store_true', default=True)
parser.add_argument('--test', dest='train', action='store_false', default=True)
parser.add_argument('--steps', dest='steps', action='store', default=10000, type=int)
parser.add_argument('--visualize', dest='visualize', action='store_true', default=False)
parser.add_argument('--model', dest='model', action='store', default="example.h5f")
args = parser.parse_args()

# set to get observation in array
#def _new_step(self, action, project=True, obs_as_dict=False):
#    return super(Arm2DEnv, self).step(action, project=project, obs_as_dict=obs_as_dict)
#Arm2DEnv.step = _new_step
# Load walking environment

env = Arm2DVecEnv(args.visualize)