53 lines
1022 B
Python
53 lines
1022 B
Python
# %%
|
|
from sb3_contrib import TRPO
|
|
from intersim.envs import IntersimpleLidarFlat
|
|
from intersim.envs.intersimple import speed_reward
|
|
import functools
|
|
|
|
model_name = "trpo_speed_lidar"
|
|
|
|
#def reward(state, action, info):
|
|
# speed = state[2].item()
|
|
# r = speed if speed < 10 else (10 - 5 * (speed - 10))
|
|
# return 0.1 * r
|
|
|
|
env = IntersimpleLidarFlat(
|
|
n_rays=5,
|
|
agent=51,
|
|
reward=functools.partial(
|
|
speed_reward,
|
|
collision_penalty=0
|
|
),
|
|
)
|
|
|
|
# %%
|
|
model = TRPO(
|
|
"MlpPolicy", env,
|
|
learning_rate=1e-4,
|
|
verbose=1,
|
|
tensorboard_log='runs/',
|
|
#use_sde=True,
|
|
#sde_sample_freq=4,
|
|
)
|
|
model.learn(total_timesteps=1000000)
|
|
model.save(model_name)
|
|
|
|
print('Done training.')
|
|
|
|
del model # remove to demonstrate saving and loading
|
|
|
|
# %%
|
|
model = TRPO.load(model_name)
|
|
|
|
obs = env.reset()
|
|
while True:
|
|
action, _states = model.predict(obs)
|
|
obs, rewards, done, info = env.step(action)
|
|
env.render(mode='post')
|
|
if done:
|
|
break
|
|
|
|
env.close(filestr='render/'+model_name)
|
|
|
|
# %%
|