Plot decision tree in R Caret

R programming

data visualization

decision tree

caret package

machine learning

Plot decision tree in R Caret

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Machine Learning Data Science Data Structures

Introduction

The caret package in R provides a unified interface for training machine learning models, including decision trees. To plot a decision tree trained with caret, you extract the underlying model object and pass it to a visualization function. For rpart trees, use rpart.plot() from the rpart.plot package. For party trees, use plot() directly. Caret wraps the model fitting but does not provide its own plotting — you always work with the underlying model's plot functions.

Training a Decision Tree with Caret

1library(caret)
2library(rpart)
3library(rpart.plot)
4
5# Use the iris dataset
6data(iris)
7
8# Train a decision tree using rpart method
9set.seed(42)
10model <- train(
11  Species ~ .,
12  data = iris,
13  method = "rpart",
14  trControl = trainControl(method = "cv", number = 10),
15  tuneLength = 5
16)
17
18print(model)
19# Shows the best cp (complexity parameter) selected by cross-validation

Plotting with rpart.plot

1# Extract the underlying rpart model
2tree <- model$finalModel
3
4# Basic plot
5rpart.plot(tree)
6
7# Enhanced plot with more detail
8rpart.plot(
9  tree,
10  type = 4,         # Label all nodes
11  extra = 104,      # Show class probability and percentage
12  fallen.leaves = TRUE,
13  main = "Iris Species Classification Tree",
14  box.palette = "RdYlGn"
15)

The type and extra parameters control how much information appears on each node:

`type`	Description
0	Draw split labels on branches
1	Label all nodes with split info
2	Label with split and node number
3	Label with split condition at each node
4	Like type 3 but includes class label

`extra`	Description
0	No extra info
1	Number of observations
2	Classification rate
100	Percentage of observations
104	Class probability and percentage

Using the Built-in plot Function

1# Base R plotting (less informative than rpart.plot)
2plot(tree)
3text(tree, use.n = TRUE, cex = 0.8)
4
5# Better: use rpart.plot for publication-quality trees
6rpart.plot(tree, roundint = FALSE, digits = 3)

Regression Tree Example

1library(caret)
2library(rpart.plot)
3
4# Regression tree with mtcars
5set.seed(42)
6reg_model <- train(
7  mpg ~ .,
8  data = mtcars,
9  method = "rpart",
10  trControl = trainControl(method = "cv", number = 5)
11)
12
13# Plot
14rpart.plot(
15  reg_model$finalModel,
16  type = 4,
17  extra = 101,    # Show n and percentage
18  digits = 2,
19  main = "MPG Prediction Tree"
20)

For regression trees, each leaf shows the predicted value (mean of observations in that node) instead of a class label.

Tuning and Plotting with Different Complexity

1# Train with explicit tuning grid
2grid <- expand.grid(cp = seq(0.001, 0.1, by = 0.01))
3
4model <- train(
5  Species ~ .,
6  data = iris,
7  method = "rpart",
8  trControl = trainControl(method = "cv", number = 10),
9  tuneGrid = grid
10)
11
12# Plot cross-validation accuracy vs complexity parameter
13plot(model)
14
15# Plot the final tree
16rpart.plot(model$finalModel, main = paste("Best cp =", model$bestTune$cp))

Lower cp values produce more complex trees (more splits). Cross-validation helps find the optimal cp that balances accuracy and simplicity.

Using party/ctree for Conditional Trees

1library(caret)
2library(party)
3
4# Train using ctree method
5model_ctree <- train(
6  Species ~ .,
7  data = iris,
8  method = "ctree",
9  trControl = trainControl(method = "cv", number = 10)
10)
11
12# Plot — party trees have their own plot method
13plot(model_ctree$finalModel)
14
15# Customized plot
16plot(model_ctree$finalModel,
17     main = "Conditional Inference Tree",
18     tp_args = list(fill = c("red", "green", "blue")))

ctree uses statistical tests (p-values) to determine splits, producing different trees than rpart.

Saving the Plot

1# Save as PNG
2png("decision_tree.png", width = 1200, height = 800, res = 150)
3rpart.plot(model$finalModel, type = 4, extra = 104)
4dev.off()
5
6# Save as PDF (vector format)
7pdf("decision_tree.pdf", width = 10, height = 7)
8rpart.plot(model$finalModel, type = 4, extra = 104)
9dev.off()

Variable Importance

1# Get variable importance from the caret model
2importance <- varImp(model)
3print(importance)
4plot(importance, main = "Variable Importance")
5
6# From the rpart model directly
7tree <- model$finalModel
8barplot(tree$variable.importance,
9        main = "Variable Importance",
10        las = 2, col = "steelblue")

Common Pitfalls

Accessing model instead of model$finalModel: Caret's train() returns a caret model object, not an rpart object. Passing model directly to rpart.plot() fails. Always extract model$finalModel for the underlying rpart tree.
Not installing rpart.plot separately: The rpart package only provides the basic plot() function for trees. The rpart.plot package (a separate CRAN package) is needed for the detailed, publication-quality plots with rpart.plot().
Tree is a single root node with no splits: If the complexity parameter cp is too high, the tree is pruned to a single node. Lower the cp value in the tuning grid or set tuneLength to a larger number to explore more options.
Using rpart.plot with non-rpart models: rpart.plot() only works with rpart objects. Trees from ctree, randomForest, or gbm require their own visualization functions. Check model$method to determine the model type.
Forgetting set.seed() before training: Decision tree training with cross-validation is stochastic. Without set.seed(), you get different trees on each run, making results non-reproducible.

Summary

Train a decision tree with caret::train(method = "rpart") and extract the tree with model$finalModel
Plot with rpart.plot(model$finalModel, type = 4, extra = 104) for detailed, readable trees
Use type and extra parameters to control the level of detail shown on nodes
For conditional inference trees, use method = "ctree" and plot with plot(model$finalModel)
Use varImp(model) to visualize which features the tree considers most important
Always set set.seed() before training for reproducible results

DSA Fundamentals

Master algorithmic patterns and data structures through hands-on LeetCode-style problems - from arrays and hashing to dynamic programming and advanced graphs.

View the course

Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Plot decision tree in R Caret

ML System Design practice on Codemia

Introduction

Training a Decision Tree with Caret

Plotting with rpart.plot

Using the Built-in plot Function

Regression Tree Example

Tuning and Plotting with Different Complexity

Using party/ctree for Conditional Trees

Saving the Plot

Variable Importance

Common Pitfalls

Summary

Related reading

DSA Fundamentals

Track what you have practised

ML System Design practice on Codemia