diff --git a/R/robin_glm.R b/R/robin_glm.R
index e604ba5..050b594 100644
--- a/R/robin_glm.R
+++ b/R/robin_glm.R
@@ -10,8 +10,10 @@
 #' @param vcov (`function`) A function to calculate the variance-covariance matrix of the treatment effect,
 #' including `vcovHC` and `gvcov`.
 #' @param family (`family`) A family object of the glm model.
-#' @param ... Additional arguments passed to `vcov`. For finer control of glm, refer to usage of `treatment_effect`,
-#' `difference`, `risk_ratio`, `odds_ratio`.
+#' @param vcov_args (`list`) Additional arguments passed to `vcov`.
+#' @param ... Additional arguments passed to `glm` or `glm.nb`.
+#' @details
+#' If family is `MASS::negative.binomial(NA)`, the function will use `MASS::glm.nb` instead of `glm`.
 #' @export
 #' @examples
 #' robin_glm(
@@ -21,14 +23,23 @@
 #' )
 robin_glm <- function(
     formula, data, treatment, contrast = "difference",
-    contrast_jac = NULL, vcov = gvcov, family = gaussian, ...) {
+    contrast_jac = NULL, vcov = "gvcov", family = gaussian(), vcov_args = list(), ...) {
   attr(formula, ".Environment") <- environment()
-  fit <- glm(formula, family = family, data = data)
+  # check if using negative.binomial family with NA as theta.
+  # If so, use MASS::glm.nb instead of glm.
+  if (identical(family$family, "Negative Binomial(NA)")) {
+    fit <- MASS::glm.nb(formula, data = data, ...)
+  } else {
+    fit <- glm(formula, family = family, data = data, ...)
+  }
   pc <- predict_counterfactual(fit, treatment, data)
   has_interaction <- h_interaction(formula, treatment)
-  if (has_interaction && identical(vcov, vcovHC) && !identical(contrast, "difference")) {
+  if (
+    has_interaction && (identical(vcov, "vcovHC") || identical(vcov, vcovHC)) &&
+      !identical(contrast, "difference")) {
     stop(
-      "Huber-White variance estimator is ONLY supported when the expected outcome difference is estimated using a linear model without treatment-covariate interactions; see the 2023 FDA guidance."
+      "Huber-White variance estimator is ONLY supported when the expected outcome difference is estimated",
+      "using a linear model without treatment-covariate interactions; see the 2023 FDA guidance."
     )
   }
   if (identical(contrast, "difference")) {
@@ -45,7 +56,7 @@ robin_glm <- function(
         numDeriv::jacobian(contrast, x)
       }
     }
-    treatment_effect(pc, eff_measure = contrast, eff_jacobian = contrast_jac, variance = vcov, ...)
+    treatment_effect(pc, eff_measure = contrast, eff_jacobian = contrast_jac, variance = vcov, vcov_args = vcov_args)
   }
 }
 
diff --git a/R/treatment_effect.R b/R/treatment_effect.R
index 0405dd2..92f7415 100644
--- a/R/treatment_effect.R
+++ b/R/treatment_effect.R
@@ -6,21 +6,26 @@
 #' @param variance (`function`) Variance function.
 #' @param eff_measure (`function`) Treatment effect measurement function.
 #' @param eff_jacobian (`function`) Treatment effect jacobian function.
+#' @param vcov_args (`list`) Additional arguments for variance.
 #' @param ... Additional arguments for variance.
 #'
 #' @export
-treatment_effect <- function(object, pair, variance, eff_measure, eff_jacobian, ...) {
+treatment_effect <- function(object, pair, variance, eff_measure, eff_jacobian, vcov_args, ...) {
   UseMethod("treatment_effect")
 }
 
 #' @export
 treatment_effect.prediction_cf <- function(
-    object, pair = names(object), variance = gvcov, eff_measure, eff_jacobian, ...) {
+    object, pair = names(object), variance = "gvcov", eff_measure, eff_jacobian, vcov_args = list(), ...) {
   assert(
+    test_string(variance),
     test_function(variance),
     test_null(variance)
   )
   assert_function(eff_measure)
+  if (missing(pair)) {
+    pair <- names(object)
+  }
   assert_vector(pair)
   assert(
     test_subset(pair, names(object)),
@@ -30,8 +35,16 @@ treatment_effect.prediction_cf <- function(
     pair <- names(object)[pair]
   }
   trt_effect <- unname(eff_measure(object[pair]))
+  if (test_string(variance)) {
+    variance_name <- variance
+    variance <- match.fun(variance)
+  } else if (test_function(variance)) {
+    variance_name <- "function"
+  } else {
+    variance_name <- "none"
+  }
   if (!is.null(variance)) {
-    inner_variance <- variance(object, ...)[pair, pair]
+    inner_variance <- do.call(variance, c(list(object), vcov_args))[pair, pair]
     if (missing(eff_jacobian)) {
       trt_jac <- numDeriv::jacobian(eff_measure, object[pair])
     } else {
@@ -47,8 +60,9 @@ treatment_effect.prediction_cf <- function(
   structure(
     .Data = trt_effect,
     name = pair_names[lower.tri(pair_names)],
+    marginal_mean = object,
     fit = attr(object, "fit"),
-    vartype = deparse(substitute(variance)),
+    vartype = variance_name,
     treatment = attr(object, "treatment_formula"),
     variance = diag(trt_var),
     class = "treatment_effect"
@@ -59,26 +73,18 @@ treatment_effect.prediction_cf <- function(
 #' @export
 #' @inheritParams predict_counterfactual
 treatment_effect.lm <- function(
-    object, pair = names(object), variance = gvcov, eff_measure, eff_jacobian,
-    treatment, data = find_data(object), ...) {
+    object, pair, variance = gvcov, eff_measure, eff_jacobian,
+    treatment, vcov_args = list(), data = find_data(object), ...) {
   pc <- predict_counterfactual(object, data = data, treatment)
-  if (missing(pair)) {
-    treatment_effect(pc, pair = , , variance = variance, eff_measure = eff_measure, eff_jacobian = eff_jacobian, ...)
-  } else {
-    treatment_effect(pc, pair, variance = variance, eff_measure = eff_measure, eff_jacobian = eff_jacobian, ...)
-  }
+  treatment_effect(pc, pair = pair, variance = variance, eff_measure = eff_measure, eff_jacobian = eff_jacobian, ...)
 }
 
 #' @export
 treatment_effect.glm <- function(
     object, pair, variance = gvcov, eff_measure, eff_jacobian,
-    treatment, data = find_data(object), ...) {
+    treatment, vcov_args = list(), data = find_data(object), ...) {
   pc <- predict_counterfactual(object, treatment, data)
-  if (missing(pair)) {
-    treatment_effect(pc, pair = , , variance = variance, eff_measure = eff_measure, eff_jacobian = eff_jacobian, ...)
-  } else {
-    treatment_effect(pc, pair, variance = variance, eff_measure = eff_measure, eff_jacobian = eff_jacobian, ...)
-  }
+  treatment_effect(pc, pair = pair, variance = variance, eff_measure = eff_measure, eff_jacobian = eff_jacobian, ...)
 }
 
 #' @rdname treatment_effect
@@ -173,11 +179,18 @@ h_lower_tri_idx <- function(n) {
 print.treatment_effect <- function(x, ...) {
   cat("Treatment Effect\n")
   cat("-------------\n")
-  cat("Model        : ", deparse(attr(x, "fit")$formula), "\n")
+  cat("Model        : ", deparse(as.formula(attr(x, "fit"))), "\n")
   cat("Randomization: ", deparse(attr(x, "treatment")), "\n")
+  cat("Marginal Mean: \n")
+  print(attr(x, "marginal_mean"))
+  
   cat("Variance Type: ", attr(x, "vartype"), "\n")
-  trt_sd <- sqrt(attr(x, "variance"))
-  z_value <- x / trt_sd
+  if (identical(attr(x, "vartype"), "none")) {
+    trt_sd <- rep(NA, length(x))
+  } else {
+    trt_sd <- sqrt(attr(x, "variance"))
+  }
+  z_value <- as.numeric(x) / trt_sd
   p <- 2 * pnorm(abs(z_value), lower.tail = FALSE)
   coef_mat <- matrix(
     c(
@@ -194,5 +207,5 @@ print.treatment_effect <- function(x, ...) {
     coef_mat,
     zap.ind = 3,
     digits = 3
-  )
+  )  
 }
diff --git a/man/robin_glm.Rd b/man/robin_glm.Rd
index 01271b9..21e697b 100644
--- a/man/robin_glm.Rd
+++ b/man/robin_glm.Rd
@@ -10,8 +10,9 @@ robin_glm(
   treatment,
   contrast = "difference",
   contrast_jac = NULL,
-  vcov = gvcov,
-  family = gaussian,
+  vcov = "gvcov",
+  family = gaussian(),
+  vcov_args = list(),
   ...
 )
 }
@@ -34,12 +35,16 @@ including \code{vcovHC} and \code{gvcov}.}
 
 \item{family}{(\code{family}) A family object of the glm model.}
 
-\item{...}{Additional arguments passed to \code{vcov}. For finer control of glm, refer to usage of \code{treatment_effect},
-\code{difference}, \code{risk_ratio}, \code{odds_ratio}.}
+\item{vcov_args}{(\code{list}) Additional arguments passed to \code{vcov}.}
+
+\item{...}{Additional arguments passed to \code{glm} or \code{glm.nb}.}
 }
 \description{
 Covariate adjusted glm model
 }
+\details{
+If family is \code{MASS::negative.binomial(NA)}, the function will use \code{MASS::glm.nb} instead of \code{glm}.
+}
 \examples{
 robin_glm(
   y ~ treatment * s1,
diff --git a/man/treatment_effect.Rd b/man/treatment_effect.Rd
index fc5b44a..1f9827e 100644
--- a/man/treatment_effect.Rd
+++ b/man/treatment_effect.Rd
@@ -7,7 +7,15 @@
 \alias{odds_ratio}
 \title{Treatment Effect}
 \usage{
-treatment_effect(object, pair, variance, eff_measure, eff_jacobian, ...)
+treatment_effect(
+  object,
+  pair,
+  variance,
+  eff_measure,
+  eff_jacobian,
+  vcov_args,
+  ...
+)
 
 difference(object, ...)
 
@@ -26,6 +34,8 @@ odds_ratio(object, ...)
 
 \item{eff_jacobian}{(\code{function}) Treatment effect jacobian function.}
 
+\item{vcov_args}{(\code{list}) Additional arguments for variance.}
+
 \item{...}{Additional arguments for variance.}
 }
 \description{
diff --git a/tests/testthat/_snaps/treatment_effect.md b/tests/testthat/_snaps/treatment_effect.md
index 4ce0541..ff3d830 100644
--- a/tests/testthat/_snaps/treatment_effect.md
+++ b/tests/testthat/_snaps/treatment_effect.md
@@ -7,7 +7,13 @@
       -------------
       Model        :  y_b ~ treatment * s1 + covar 
       Randomization:  treatment ~ s1 
-      Variance Type:  variance 
+      Marginal Mean: 
+      counter-factual prediction
+      
+            pbo      trt1      trt2 
+      0.3560965 0.5806957 0.6213865 
+      
+      Variance Type:  function 
                   Estimate Std.Err Z Value  Pr(>z)    
       trt1 - pbo    0.2246  0.0477    4.71 2.5e-06 ***
       trt2 - pbo    0.2653  0.0475    5.58 2.4e-08 ***
@@ -22,9 +28,15 @@
     Output
       Treatment Effect
       -------------
-      Model        :  NULL 
+      Model        :  y ~ treatment * s1 + covar 
       Randomization:  treatment ~ s1 
-      Variance Type:  variance 
+      Marginal Mean: 
+      counter-factual prediction
+      
+            pbo      trt1      trt2 
+      0.2003208 0.7639709 0.9712499 
+      
+      Variance Type:  function 
                   Estimate Std.Err Z Value  Pr(>z)    
       trt1 - pbo     0.564   0.101    5.60 2.2e-08 ***
       trt2 - pbo     0.771   0.101    7.61 2.8e-14 ***
@@ -32,3 +44,47 @@
       ---
       Signif. codes:  0 '***' 0.001 '**' 0.01 '*' 0.05 '.' 0.1 ' ' 1
 
+# treatment_effect works if variance is not used
+
+    Code
+      treatment_effect(fit_binom, treatment = treatment ~ s1, eff_measure = h_diff,
+      variance = NULL)
+    Output
+      Treatment Effect
+      -------------
+      Model        :  y_b ~ treatment * s1 + covar 
+      Randomization:  treatment ~ s1 
+      Marginal Mean: 
+      counter-factual prediction
+      
+            pbo      trt1      trt2 
+      0.3560965 0.5806957 0.6213865 
+      
+      Variance Type:  none 
+                  Estimate Std.Err Z Value Pr(>z)
+      trt1 - pbo    0.2246      NA      NA     NA
+      trt2 - pbo    0.2653      NA      NA     NA
+      trt2 - trt1   0.0407      NA      NA     NA
+
+# treatment_effect works if pair is integer
+
+    Code
+      treatment_effect(fit_binom, pair = c(1, 2), treatment = treatment ~ s1,
+      eff_measure = h_diff)
+    Output
+      Treatment Effect
+      -------------
+      Model        :  y_b ~ treatment * s1 + covar 
+      Randomization:  treatment ~ s1 
+      Marginal Mean: 
+      counter-factual prediction
+      
+            pbo      trt1      trt2 
+      0.3560965 0.5806957 0.6213865 
+      
+      Variance Type:  function 
+                 Estimate Std.Err Z Value  Pr(>z)    
+      trt1 - pbo   0.2246  0.0477    4.71 2.5e-06 ***
+      ---
+      Signif. codes:  0 '***' 0.001 '**' 0.01 '*' 0.05 '.' 0.1 ' ' 1
+
diff --git a/tests/testthat/_snaps/variance.md b/tests/testthat/_snaps/variance.md
index c3eeba1..ef8bb18 100644
--- a/tests/testthat/_snaps/variance.md
+++ b/tests/testthat/_snaps/variance.md
@@ -8,6 +8,16 @@
       trt1  4.523445e-07  1.164889e-03 -7.709031e-07
       trt2 -9.709004e-06 -7.709031e-07  1.170214e-03
 
+---
+
+    Code
+      vcovHC(pc)
+    Output
+                     pbo          trt1          trt2
+      pbo   1.106024e-03  4.523445e-07 -9.709004e-06
+      trt1  4.523445e-07  1.164889e-03 -7.709031e-07
+      trt2 -9.709004e-06 -7.709031e-07  1.170214e-03
+
 # gvcov works
 
     Code
@@ -21,7 +31,37 @@
 ---
 
     Code
-      gvcov(pc, randomization = "permute_block")
+      gvcov(pc)
+    Output
+                    pbo         trt1         trt2
+      pbo  1.128902e-03 1.856234e-05 1.333885e-05
+      trt1 1.856234e-05 1.184599e-03 2.178112e-05
+      trt2 1.333885e-05 2.178112e-05 1.157268e-03
+
+---
+
+    Code
+      gvcov(pc)
+    Output
+                    pbo         trt1         trt2
+      pbo  1.128902e-03 1.856234e-05 1.333885e-05
+      trt1 1.856234e-05 1.184599e-03 2.178112e-05
+      trt2 1.333885e-05 2.178112e-05 1.157268e-03
+
+---
+
+    Code
+      gvcov(pc, decompose = FALSE)
+    Output
+                    pbo         trt1         trt2
+      pbo  1.127076e-03 1.856234e-05 1.333885e-05
+      trt1 1.856234e-05 1.179430e-03 2.178112e-05
+      trt2 1.333885e-05 2.178112e-05 1.164046e-03
+
+---
+
+    Code
+      gvcov(pc)
     Output
                     pbo         trt1         trt2
       pbo  1.128902e-03 1.856234e-05 1.333885e-05
diff --git a/tests/testthat/test-robin_glm.R b/tests/testthat/test-robin_glm.R
index 4a124b8..4e640a7 100644
--- a/tests/testthat/test-robin_glm.R
+++ b/tests/testthat/test-robin_glm.R
@@ -14,7 +14,7 @@ test_that("robin_glm works correctly", {
     robin_glm(
       y ~ treatment * s1,
       data = dummy_data, treatment = treatment ~ s1,
-      contrast = "difference", vcov = vcovHC
+      contrast = "difference", vcov = "vcovHC"
     )
   )
   expect_silent(robin_glm(y_b ~ treatment * s1, data = dummy_data, treatment = treatment ~ s1, contrast = "difference"))
@@ -24,9 +24,23 @@ test_that("robin_glm works correctly", {
     robin_glm(
       y_b ~ treatment * s1,
       data = dummy_data, treatment = treatment ~ s1,
-      contrast = "odds_ratio", vcov = vcovHC
+      contrast = "odds_ratio", vcov = "vcovHC"
     ),
     "Huber-White variance estimator is ONLY"
   )
   expect_silent(robin_glm(y_b ~ treatment * s1, data = dummy_data, treatment = treatment ~ s1, contrast = h_diff))
 })
+
+
+test_that("robin_glm works for glm.nb", {
+  dummy_data2 <- dummy_data
+  dummy_data2$y_b <- rep(seq_len(10), nrow(dummy_data2) / 10)
+  expect_silent(
+    robin_glm(
+      y_b ~ treatment * s1,
+      data = dummy_data2, treatment = treatment ~ s1,
+      family = MASS::negative.binomial(theta = NA),
+      contrast = "difference", vcov = "vcovHC"
+    )
+  )
+})
diff --git a/tests/testthat/test-treatment_effect.R b/tests/testthat/test-treatment_effect.R
index 987e4a0..8a545e1 100644
--- a/tests/testthat/test-treatment_effect.R
+++ b/tests/testthat/test-treatment_effect.R
@@ -149,3 +149,11 @@ test_that("treatment_effect works for lm/glm object", {
   expect_snapshot(treatment_effect(fit_binom, treatment = treatment ~ s1, eff_measure = h_diff))
   expect_snapshot(treatment_effect(fit_lm, treatment = treatment ~ s1, eff_measure = h_diff))
 })
+
+test_that("treatment_effect works if variance is not used", {
+  expect_snapshot(treatment_effect(fit_binom, treatment = treatment ~ s1, eff_measure = h_diff, variance = NULL))
+})
+
+test_that("treatment_effect works if pair is integer", {
+  expect_snapshot(treatment_effect(fit_binom, pair = c(1, 2), treatment = treatment ~ s1, eff_measure = h_diff))
+})
diff --git a/tests/testthat/test-utils.R b/tests/testthat/test-utils.R
index 92a29f5..6da6ffd 100644
--- a/tests/testthat/test-utils.R
+++ b/tests/testthat/test-utils.R
@@ -6,6 +6,14 @@ test_that("h_get_vars works for formula", {
     h_get_vars("treatment"),
     "Must be a formula, not character"
   )
+  expect_error(
+    h_get_vars(trt ~ pb(s1) + ps(s2)),
+    "only one randomization schema is allowed!"
+  )
+  expect_error(
+    h_get_vars( ~ ps(s2)),
+    "treatment formula must be of type treatment ~ strata"
+  )
   expect_error(
     h_get_vars(NULL),
     "Must be a formula, not 'NULL'"
diff --git a/tests/testthat/test-variance.R b/tests/testthat/test-variance.R
index a6d5136..3a81b8f 100644
--- a/tests/testthat/test-variance.R
+++ b/tests/testthat/test-variance.R
@@ -3,6 +3,10 @@ test_that("vcovHC works", {
   expect_snapshot(
     vcovHC(pc)
   )
+  pc <- predict_counterfactual(fit_binom, treatment ~ pb(s1))
+  expect_snapshot(
+    vcovHC(pc)
+  )
 })
 
 test_that("gvcov works", {
@@ -10,7 +14,19 @@ test_that("gvcov works", {
   expect_snapshot(
     gvcov(pc)
   )
+  pc <- predict_counterfactual(fit_binom, treatment ~ 1)
+  expect_snapshot(
+    gvcov(pc)
+  )
+  pc <- predict_counterfactual(fit_binom, treatment ~ pb(s1))
   expect_snapshot(
-    gvcov(pc, randomization = "permute_block")
+    gvcov(pc)
+  )
+  expect_snapshot(
+    gvcov(pc, decompose = FALSE)
+  )
+  pc <- predict_counterfactual(fit_binom, treatment ~ ps(s1))
+  expect_snapshot(
+    gvcov(pc)
   )
 })