Create a dataset using an agent (#306)

Generate a dataset using an agent (allowing to select between this and a random dataset)

rl_coach/agents/qr_dqn_agent.py

@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-
+from copy import copy
 from typing import Union
 
 import numpy as np
@@ -79,6 +79,17 @@ class QuantileRegressionDQNAgent(ValueOptimizationAgent):
             actions_q_values = None
         return actions_q_values
 
+    # prediction's format is (batch,actions,atoms)
+    def get_all_q_values_for_states_and_softmax_probabilities(self, states: StateType):
+        actions_q_values, softmax_probabilities = None, None
+        if self.exploration_policy.requires_action_values():
+            outputs = copy(self.networks['main'].online_network.outputs)
+            outputs.append(self.networks['main'].online_network.output_heads[0].softmax)
+            quantile_values, softmax_probabilities = self.get_prediction(states, outputs)
+            actions_q_values = self.get_q_values(quantile_values)
+
+        return actions_q_values, softmax_probabilities
+
     def learn_from_batch(self, batch):
         network_keys = self.ap.network_wrappers['main'].input_embedders_parameters.keys()