79 lines
2.2 KiB
Python
79 lines
2.2 KiB
Python
# %%
|
|
import torch
|
|
from core.policy import SetPolicy
|
|
from tqdm import tqdm
|
|
|
|
expert_data = torch.load('intersimple-expert-data-setobs2.pt')
|
|
states, actions, _, dones = expert_data
|
|
|
|
policy = SetPolicy(actions.shape[-1])
|
|
|
|
policy = policy.cuda()
|
|
optim = torch.optim.Adam(policy.parameters(), lr=1e-4)
|
|
states = states[~dones].cuda()
|
|
actions = actions[~dones].cuda()
|
|
|
|
for _ in tqdm(range(10000)):
|
|
optim.zero_grad()
|
|
loss = -policy.log_prob(policy(states), actions).mean()
|
|
loss.backward()
|
|
optim.step()
|
|
|
|
print('Loss', loss)
|
|
|
|
torch.save(policy.state_dict(), 'bc-intersimple-setobs2.pt')
|
|
|
|
# %%
|
|
import numpy as np
|
|
from core.policy import SetPolicy
|
|
from util.wrappers import Setobs, TransformObservation, CollisionPenaltyWrapper
|
|
from intersim.envs import IntersimpleLidarFlatRandom
|
|
from intersim.envs.intersimple import speed_reward
|
|
import functools
|
|
|
|
policy = SetPolicy(actions.shape[-1])
|
|
policy.load_state_dict(torch.load('bc-intersimple-setobs2.pt'))
|
|
|
|
obs_min = np.array([
|
|
[-1000, -1000, 0, -np.pi, -1e-1, 0.],
|
|
[0, -np.pi, -20, -20, -np.pi, -1e-1],
|
|
[0, -np.pi, -20, -20, -np.pi, -1e-1],
|
|
[0, -np.pi, -20, -20, -np.pi, -1e-1],
|
|
[0, -np.pi, -20, -20, -np.pi, -1e-1],
|
|
[0, -np.pi, -20, -20, -np.pi, -1e-1],
|
|
]).reshape(-1)
|
|
|
|
obs_max = np.array([
|
|
[1000, 1000, 20, np.pi, 1e-1, 0.],
|
|
[50, np.pi, 20, 20, np.pi, 1e-1],
|
|
[50, np.pi, 20, 20, np.pi, 1e-1],
|
|
[50, np.pi, 20, 20, np.pi, 1e-1],
|
|
[50, np.pi, 20, 20, np.pi, 1e-1],
|
|
[50, np.pi, 20, 20, np.pi, 1e-1],
|
|
]).reshape(-1)
|
|
|
|
env = Setobs(
|
|
TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlatRandom(
|
|
n_rays=5,
|
|
reward=functools.partial(
|
|
speed_reward,
|
|
collision_penalty=0
|
|
),
|
|
stop_on_collision=False,
|
|
), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))
|
|
)
|
|
|
|
obs = env.reset()
|
|
env.render(mode='post')
|
|
for i in range(300):
|
|
#action, _ = policy.predict(torch.tensor(obs))
|
|
action = policy.sample(policy(torch.tensor(obs, dtype=torch.float32)))
|
|
obs, reward, done, _ = env.step(action)
|
|
env.render(mode='post')
|
|
print('step', i, 'reward', reward)
|
|
if done:
|
|
break
|
|
env.close()
|
|
|
|
# %%
|