Add DirichletFullyConnectedActor to Soft Actor-Critic

Summary: This can be used for problems where the action vector must sum to 1

Reviewed By: kittipatv

Differential Revision: D15206348

fbshipit-source-id: 665fbed893d8c52d451a12d3bb2e73b2638b7963

Author: econti

test/onnx/test_pytorch_onnx_caffe2.py

@@ -1606,6 +1606,19 @@ def forward(self, input):
         x = torch.empty(BATCH_SIZE, 10, 10).uniform_(4, 9)
         self.run_model_test(Log2Model(), train=False, input=x, batch_size=BATCH_SIZE)
 
+    def test_sample_dirichlet(self):
+        class DirichletModel(torch.nn.Module):
+            def forward(self, input):
+                return torch._sample_dirichlet(input)
+
+        x = torch.randn(2, 3, 4, requires_grad=False)
+        model = DirichletModel()
+        onnxir, _ = do_export(model, x)
+        onnx_model = onnx.ModelProto.FromString(onnxir)
+        prepared = c2.prepare(onnx_model)
+        caffe2_out = prepared.run(inputs=[x.cpu().numpy()])
+        self.assertEqual(caffe2_out[0].shape, x.shape)
+
     def test_prim_shape(self):
         x = torch.randn(4, 5, requires_grad=True)
         @torch.jit.script

torch/onnx/symbolic.py

@@ -423,6 +423,13 @@ def cumsum(g, input, dim):
     return g.op("ATen", input, operator_s="cumsum", dim_i=dim)
 
 
+def _sample_dirichlet(g, self, generator):
+    if not generator.node().mustBeNone():
+        return _unimplemented('_sample_dirichlet',
+                              'We are not able to export generator')
+    return g.op("ATen", self, operator_s="_sample_dirichlet")
+
+
 def t(g, self):
     return g.op("Transpose", self, perm_i=(1, 0))