Add axes to action function image

Author: wedesoft

src/ppo/main.clj

@@ -1059,20 +1059,24 @@
 (def actor (Actor 3 64 1))
 (py. actor load_state_dict (torch/load "src/ppo/actor.pt"))
 
-(let [angle-values (torch/linspace (- PI) PI 854)
-      speed-values (torch/linspace 1.0 -1.0 480)
-      grid         (torch/meshgrid speed-values angle-values :indexing "ij")
-      cos-angle    (torch/cos (last grid))
-      sin-angle    (torch/sin (last grid))
-      observations (torch/stack [(py. cos-angle ravel)
-                                 (py. sin-angle ravel)
-                                 (py. (first grid) ravel)]
-                                :axis 1)
-      actions      (without-gradient
-                     (py. (py. (py. actor deterministic_act observations)
-                               reshape 480 854) numpy))
-      actions-tensor (dtype/elemwise-cast (dtt/ensure-tensor (py/->jvm actions))
-                                          :float32)]
+(let [angle-values   (torch/linspace (- PI) PI 854)
+      speed-values   (torch/linspace 1.0 -1.0 480)
+      grid           (torch/meshgrid speed-values angle-values :indexing "ij")
+      cos-angle      (torch/cos (last grid))
+      sin-angle      (torch/sin (last grid))
+      observations   (torch/stack [(py. cos-angle ravel)
+                                   (py. sin-angle ravel)
+                                   (py. (first grid) ravel)]
+                                  :axis 1)
+      actions        (without-gradient
+                       (py. (py. (py. actor deterministic_act observations)
+                                 reshape 480 854) numpy))
+      actions-tensor (dtt/clone
+                       (dtype/elemwise-cast (dtt/ensure-tensor (py/->jvm actions))
+                                            :float32))
+      actions-trsps  (dtt/transpose actions-tensor [1 0])]
+  (dtt/mset! actions-tensor 240 (dfn/- 1.0 (actions-tensor 240)))
+  (dtt/mset! actions-trsps 427 (dfn/- 1.0 (actions-trsps 427)))
   (bufimg/tensor->image (dfn/* actions-tensor 255)))
 ;; This image shows the motor control input as a function of pendulum angle and angular velocity.
 ;; As one can see, the pendulum is decelerated when the speed is high (dark values at the top of the image).