from env import ArmEnv
from rl import DDPG


arms = 3
env = ArmEnv(n_arms=arms, random_goal=True)
model = DDPG(env.action_dim, env.state_dim, env.action_bound)
for ep in range(701):
    state = env.reset()
    for t in range(100):
        env.render()
        action = model.act(state)
        next_state, reward, done = env.step(action)
        print("{}回合，{}步，奖励={}, {}".format(ep, t, reward, "完成目标" if done else ""))
        model.store_transition(state, action, reward, next_state)
        model.learn()
    model.save_weights("models/{}arms/ep{}".format(arms, ep))