Merge branch 'main' into idm_upgrade

scratch/etienne/trpo/experiments/intersimple-expert-rollout-setobs2.py

@@ -1,10 +1,13 @@
+import sys
+sys.path.append('../../../../')
+
 import torch
 import functools
-from core.sampling import rollout_sb3
-from intersim.envs import IntersimpleLidarFlatRandom
+from src.core.sampling import rollout_sb3
+from intersim.envs import IntersimpleLidarFlatIncrementingAgent
 from intersim.envs.intersimple import speed_reward
 from intersim.expert import NormalizedIntersimpleExpert
-from util.wrappers import CollisionPenaltyWrapper, Setobs
+from src.util.wrappers import CollisionPenaltyWrapper, Setobs
 import numpy as np
 from gym.wrappers import TransformObservation
 
@@ -26,7 +29,9 @@ obs_max = np.array([
     [50, np.pi, 20, 20, np.pi, 1e-1],
 ]).reshape(-1)
 
-env = IntersimpleLidarFlatRandom(
+env = IntersimpleLidarFlatIncrementingAgent(
+    loc=0,
+    track=4,
     n_rays=5,
     reward=functools.partial(
         speed_reward,
@@ -42,7 +47,8 @@ env = Setobs(TransformObservation(
         collision_distance=6, collision_penalty=100
     ), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10)
 ))
-expert_data = rollout_sb3(env, policy, n_episodes=2048, max_steps_per_episode=200)
+print(env.nv, 'vehicles')
+expert_data = rollout_sb3(env, policy, n_episodes=150, max_steps_per_episode=200)
 
 states, actions, rewards, dones = expert_data
 print(f'Expert mean episode length {(~dones).sum() / states.shape[0]}')
@@ -50,4 +56,4 @@ print(f'Expert mean reward per episode {rewards[~dones].sum() / states.shape[0]}
 print(f'Observation mean', states[~dones].mean(0))
 print(f'Observation std', states[~dones].std(0))
 
-torch.save(expert_data, 'intersimple-expert-data-setobs2.pt')
+torch.save(expert_data, 'intersimple-expert-data-setobs2-loc0-track4.pt')

scratch/etienne/trpo/experiments/sgail-ppo-options-setobs2.py

@@ -15,79 +15,87 @@ from src.util.wrappers import CollisionPenaltyWrapper, TransformObservation, Set
 import numpy as np
 from src.safe_options.options import SafeOptionsEnv
 from torch.utils.tensorboard import SummaryWriter
+from ray import tune
 
-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)
+def training_function(config):
+    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)
+    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)
 
-envs = [SafeOptionsEnv(Setobs(
-    TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlatRandom(
-        n_rays=5,
-        reward=functools.partial(
-            speed_reward,
-            collision_penalty=0
-        ),
-        stop_on_collision=True,
-    ), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))
-), options=[(0, 5), (1, 5), (2, 5), (4, 5), (6, 5), (8, 5), (10, 5)], safe_actions_collision_method='circle', abort_unsafe_collision_method='circle') for _ in range(60)]
+    envs = [SafeOptionsEnv(Setobs(
+        TransformObservation(CollisionPenaltyWrapper(IntersimpleLidarFlatRandom(
+            n_rays=5,
+            reward=functools.partial(
+                speed_reward,
+                collision_penalty=0
+            ),
+            stop_on_collision=True,
+        ), collision_distance=6, collision_penalty=100), lambda obs: (obs - obs_min) / (obs_max - obs_min + 1e-10))
+    ), options=[(0, 5), (1, 5), (2, 5), (4, 5), (6, 5), (8, 5), (10, 5)], safe_actions_collision_method='circle', abort_unsafe_collision_method='circle') for _ in range(60)]
 
-env_fn = lambda i: envs[i]
+    env_fn = lambda i: envs[i]
 
-policy = SetMaskedDiscretePolicy(env_fn(0).action_space.n)
-pi_opt = torch.optim.Adam(policy.parameters(), lr=3e-4)
+    policy = SetMaskedDiscretePolicy(env_fn(0).action_space.n) # config net architecture
+    pi_opt = torch.optim.Adam(policy.parameters(), lr=3e-5)  # config learning rate
+    pi_lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(pi_opt, gamma=0.98) # config lr decay
 
-value = SetValue()
-v_opt = torch.optim.Adam(value.parameters(), lr=1e-3)
+    value = SetValue() # config net architecture
+    v_opt = torch.optim.Adam(value.parameters(), lr=1e-3) # config lr
 
-discriminator = DeepsetDiscriminator()
-disc_opt = torch.optim.Adam(discriminator.parameters(), lr=1e-3, weight_decay=1e-4)
+    discriminator = DeepsetDiscriminator() # config net architecture
+    disc_opt = torch.optim.Adam(discriminator.parameters(), lr=1e-3, weight_decay=1e-3) # config lr, weight decay
 
-expert_data = torch.load('intersimple-expert-data-setobs2.pt')
-expert_data = Buffer(*expert_data)
+    expert_data = torch.load('intersimple-expert-data-setobs2.pt')
+    expert_data = Buffer(*expert_data)
 
-# %%
-def callback(epoch, value, policy):
-    if not epoch % 10:
-        torch.save(policy.state_dict(), f'sgail-ppo-options-setobs2-{epoch}.pt')
-        torch.save(value.state_dict(), f'sgail-ppo-options-setobs2-value-{epoch}.pt')
+    def callback(info):
+        tune.report(gen_mean_reward_per_episode=info['gen/mean_reward_per_episode'])
 
-value, policy = gail_ppo(
-    env_fn=env_fn,
-    expert_data=expert_data,
-    discriminator=discriminator,
-    disc_opt=disc_opt,
-    disc_iters=100,
-    policy=policy,
-    value=value,
-    v_opt=v_opt,
-    v_iters=1000,
-    epochs=200,
-    rollout_episodes=60,
-    rollout_steps=60,
-    gamma=0.99,
-    gae_lambda=0.9,
-    clip_ratio=0.2,
-    pi_opt=pi_opt,
-    pi_iters=100,
-    logger=SummaryWriter(comment='sgail-ppo-options-setobs2'),
-    callback=callback,
+    value, policy = gail_ppo(
+        env_fn=env_fn,
+        expert_data=expert_data,
+        discriminator=discriminator,
+        disc_opt=disc_opt,
+        disc_iters=100, # config
+        policy=policy,
+        value=value,
+        v_opt=v_opt,
+        v_iters=1000, # config
+        epochs=200,
+        rollout_episodes=60,
+        rollout_steps=60,
+        gamma=0.99,
+        gae_lambda=0.9,
+        clip_ratio=0.2, # config
+        pi_opt=pi_opt,
+        pi_iters=100, # config
+        logger=SummaryWriter(comment='sgail-ppo-options-setobs2'),
+        callback=callback,
+        lr_schedulers=[pi_lr_scheduler],
+    )
+
+analysis = tune.run(
+    training_function,
+    config={
+        'dummy': tune.grid_search([0.001, 0.01, 0.1]),
+    }
 )
 
-torch.save(policy.state_dict(), 'sgail-ppo-options-setobs2.pt')
+print('Best config: ', analysis.get_best_config(metric='gen_mean_reward_per_episode', mode='min'))
 
 # %%
 policy = SetMaskedDiscretePolicy(env_fn(0).action_space.n)

scratch/johannes/horner_scheme.py

@@ -0,0 +1,90 @@
+# %%
+
+import numpy as np
+import torch
+from timeit import default_timer as timer
+
+# %%
+
+def powerseries(x, deg):
+    return torch.stack([x**i for i in range(deg+1)],dim=-1)
+
+def improved_powerseries(x, deg):
+    r = torch.ones(*x.shape, deg+1, dtype=torch.float64)
+    for i in range(1,deg+1):
+        r[:, :, i] = r[:, :, i-1] * x
+    return r
+
+def horner_scheme(x, poly):
+    deg = poly.shape[-1]
+    nsteps = x.shape[-1]
+    r = poly[:, -1:].repeat(1, nsteps)
+    for i in range(2, deg+1):
+        r *= x
+        r += poly[:, -i:1-i]
+    return r
+
+# %%
+
+nv = 151
+delta = 10
+n = 20
+
+state_s = torch.rand((nv, 1))
+nan_idx = np.random.choice([True, False], 151)
+state_s[nan_idx] = np.nan
+
+# %%
+
+n_coef = 21
+xpoly = torch.rand((nv, n_coef),dtype=torch.float64)
+ypoly = torch.rand((nv, n_coef),dtype=torch.float64)
+
+ds = delta * torch.arange(1,n+1).repeat(nv,1)
+
+s = ds + state_s
+s = s.type(torch.float64)
+
+smax = s[:, 0]
+smax = smax.unsqueeze(-1)
+
+
+start = timer()
+for _ in range(100):
+    deg = xpoly.shape[-1] - 1
+    expand_sims = powerseries(s, deg) # (nv, n, deg+1)
+    # print(expand_sims.shape)
+    y = (ypoly.unsqueeze(1) * expand_sims).sum(dim=-1)
+    x = (xpoly.unsqueeze(1) * expand_sims).sum(dim=-1)
+end = timer()
+print("Powerseries: {}".format((end-start)*1))
+
+start = timer()
+for _ in range(100):
+    deg = xpoly.shape[-1] - 1
+    expand_sims = improved_powerseries(s, deg) # (nv, n, deg+1)
+    # print(expand_sims.shape)
+    yp = (ypoly.unsqueeze(1) * expand_sims).sum(dim=-1)
+    xp = (xpoly.unsqueeze(1) * expand_sims).sum(dim=-1)
+end = timer()
+print("Improved Powerseries: {}".format((end-start)*1))
+
+start = timer()
+for _ in range(100):
+    x_horner = horner_scheme(s, xpoly)
+    y_horner = horner_scheme(s, ypoly)
+end = timer()
+print("Horner: {}".format((end-start)*1))
+
+start = timer()
+for _ in range(100):
+    x_max = horner_scheme(smax, xpoly)
+    y_max = horner_scheme(smax, ypoly)
+end = timer()
+# print("Horner smax: {}".format((end-start)*1))
+
+assert np.all(np.isclose(xp,x)[~nan_idx])
+assert np.all(np.isclose(yp,y)[~nan_idx])
+assert np.all(np.isclose(x_horner,x)[~nan_idx])
+assert np.all(np.isclose(y_horner,y)[~nan_idx])
+# %%