diff --git a/gtsam/hybrid/HybridFactorGraph.cpp b/gtsam/hybrid/HybridFactorGraph.cpp
new file mode 100644
index 000000000..4238925d6
--- /dev/null
+++ b/gtsam/hybrid/HybridFactorGraph.cpp
@@ -0,0 +1,78 @@
+/* ----------------------------------------------------------------------------
+
+ * GTSAM Copyright 2010-2022, Georgia Tech Research Corporation,
+ * Atlanta, Georgia 30332-0415
+ * All Rights Reserved
+ * Authors: Frank Dellaert, et al. (see THANKS for the full author list)
+
+ * See LICENSE for the license information
+
+ * -------------------------------------------------------------------------- */
+
+/**
+ * @file   HybridFactorGraph.cpp
+ * @brief  Factor graph with utilities for hybrid factors.
+ * @author Varun Agrawal
+ * @author Frank Dellaert
+ * @date   January, 2023
+ */
+
+#include <gtsam/discrete/DecisionTreeFactor.h>
+#include <gtsam/hybrid/HybridFactorGraph.h>
+
+#include <boost/format.hpp>
+
+namespace gtsam {
+
+/* ************************************************************************* */
+DiscreteKeys HybridFactorGraph::discreteKeys() const {
+  DiscreteKeys keys;
+  for (auto& factor : factors_) {
+    if (auto p = boost::dynamic_pointer_cast<DecisionTreeFactor>(factor)) {
+      for (const DiscreteKey& key : p->discreteKeys()) {
+        keys.push_back(key);
+      }
+    }
+    if (auto p = boost::dynamic_pointer_cast<HybridFactor>(factor)) {
+      for (const DiscreteKey& key : p->discreteKeys()) {
+        keys.push_back(key);
+      }
+    }
+  }
+  return keys;
+}
+
+/* ************************************************************************* */
+KeySet HybridFactorGraph::discreteKeySet() const {
+  KeySet keys;
+  for (const DiscreteKey& k : discreteKeys()) {
+    keys.insert(k.first);
+  }
+  return keys;
+}
+
+/* ************************************************************************* */
+std::unordered_map<Key, DiscreteKey> HybridFactorGraph::discreteKeyMap() const {
+  std::unordered_map<Key, DiscreteKey> result;
+  for (const DiscreteKey& k : discreteKeys()) {
+    result[k.first] = k;
+  }
+  return result;
+}
+
+/* ************************************************************************* */
+const KeySet HybridFactorGraph::continuousKeySet() const {
+  KeySet keys;
+  for (auto& factor : factors_) {
+    if (auto p = boost::dynamic_pointer_cast<HybridFactor>(factor)) {
+      for (const Key& key : p->continuousKeys()) {
+        keys.insert(key);
+      }
+    }
+  }
+  return keys;
+}
+
+/* ************************************************************************* */
+
+}  // namespace gtsam
diff --git a/gtsam/hybrid/HybridFactorGraph.h b/gtsam/hybrid/HybridFactorGraph.h
index 8e1fa0123..4d2a113c5 100644
--- a/gtsam/hybrid/HybridFactorGraph.h
+++ b/gtsam/hybrid/HybridFactorGraph.h
@@ -19,15 +19,17 @@
 
 #pragma once
 
-#include <gtsam/discrete/DiscreteFactor.h>
-#include <gtsam/hybrid/HybridDiscreteFactor.h>
 #include <gtsam/hybrid/HybridFactor.h>
 #include <gtsam/inference/FactorGraph.h>
-#include <gtsam/inference/Ordering.h>
 
 #include <boost/format.hpp>
+#include <unordered_map>
+
 namespace gtsam {
 
+class DiscreteFactor;
+class Ordering;
+
 using SharedFactor = boost::shared_ptr<Factor>;
 
 /**
@@ -80,16 +82,6 @@ class HybridFactorGraph : public FactorGraph<Factor> {
   using Base::push_back;
   using Base::resize;
 
-  /**
-   * Add a discrete factor *pointer* to the internal discrete graph
-   * @param discreteFactor - boost::shared_ptr to the factor to add
-   */
-  template <typename FACTOR>
-  IsDiscrete<FACTOR> push_discrete(
-      const boost::shared_ptr<FACTOR>& discreteFactor) {
-    Base::push_back(boost::make_shared<HybridDiscreteFactor>(discreteFactor));
-  }
-
   /**
    * Add a discrete-continuous (Hybrid) factor *pointer* to the graph
    * @param hybridFactor - boost::shared_ptr to the factor to add
@@ -99,12 +91,10 @@ class HybridFactorGraph : public FactorGraph<Factor> {
     Base::push_back(hybridFactor);
   }
 
-  /// Construct a factor and add (shared pointer to it) to factor graph.
+  /// Construct a discrete factor and add shared pointer to the factor graph.
   template <class FACTOR, class... Args>
   IsDiscrete<FACTOR> emplace_discrete(Args&&... args) {
-    auto factor = boost::allocate_shared<FACTOR>(
-        Eigen::aligned_allocator<FACTOR>(), std::forward<Args>(args)...);
-    push_discrete(factor);
+    emplace_shared<FACTOR>(std::forward<Args>(args)...);
   }
 
   /// Construct a factor and add (shared pointer to it) to factor graph.
@@ -116,30 +106,16 @@ class HybridFactorGraph : public FactorGraph<Factor> {
   }
 
   /// Get all the discrete keys in the factor graph.
-  const KeySet discreteKeys() const {
-    KeySet discrete_keys;
-    for (auto& factor : factors_) {
-      if (auto p = boost::dynamic_pointer_cast<HybridFactor>(factor)) {
-        for (const DiscreteKey& k : p->discreteKeys()) {
-          discrete_keys.insert(k.first);
-        }
-      }
-    }
-    return discrete_keys;
-  }
+  DiscreteKeys discreteKeys() const;
+
+  /// Get all the discrete keys in the factor graph, as a set.
+  KeySet discreteKeySet() const;
+
+  /// Get a map from Key to corresponding DiscreteKey.
+  std::unordered_map<Key, DiscreteKey> discreteKeyMap() const;
 
   /// Get all the continuous keys in the factor graph.
-  const KeySet continuousKeys() const {
-    KeySet keys;
-    for (auto& factor : factors_) {
-      if (auto p = boost::dynamic_pointer_cast<HybridFactor>(factor)) {
-        for (const Key& key : p->continuousKeys()) {
-          keys.insert(key);
-        }
-      }
-    }
-    return keys;
-  }
+  const KeySet continuousKeySet() const;
 };
 
 }  // namespace gtsam