Port TRPO, PPO, GAIL

scratch/etienne/trpo/core/optimization.py

@@ -0,0 +1,39 @@
+import torch
+
+def conjugate_gradient(A, b, max_iters, res_tol=1e-10):
+    x = torch.zeros_like(b)
+    r = b - A(x)
+    p = r
+
+    rTr = r.T @ r
+
+    for _ in range(max_iters):
+        Ap = A(p)
+        alpha = rTr / (p.T @ Ap)
+        x = x + alpha * p
+
+        r = r - alpha * Ap
+        if torch.norm(r) < res_tol:
+            break
+
+        rTrnew = r.T @ r
+        beta = rTrnew / rTr
+        p = r + beta * p
+        rTr = rTrnew
+
+    return x
+
+def line_search(f, x0, dx, g0, alpha, condition, max_steps=10, c1=0.1):
+    assert 0 < alpha < 1
+    
+    f0 = f(x0)
+    for _ in range(max_steps):
+        x = x0 + dx
+
+        if (f(x) > f0 + c1 * g0.T @ dx) and condition(x):
+            return x
+        
+        dx *= alpha
+    
+    print('Line search failed, returning x0')
+    return x0